import asyncio import logging import time from collections import deque from datetime import datetime, timedelta import matplotlib.pyplot as plt import numpy as np import pandas as pd from matplotlib.colors import LogNorm from core.data_provider import DataProvider, MarketTick from core.config import get_config # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) class COBStabilityTester: def __init__(self, symbol='ETHUSDT', duration_seconds=15): self.symbol = symbol self.duration = timedelta(seconds=duration_seconds) self.ticks = deque() # Set granularity (buckets) based on symbol if 'ETH' in symbol.upper(): self.price_granularity = 1.0 # 1 USD for ETH elif 'BTC' in symbol.upper(): self.price_granularity = 10.0 # 10 USD for BTC else: self.price_granularity = 1.0 # Default 1 USD logger.info(f"Using price granularity: ${self.price_granularity} for {symbol}") # Initialize DataProvider the same way as clean_dashboard logger.info("Initializing DataProvider like in clean_dashboard...") self.data_provider = DataProvider() # Use default constructor like clean_dashboard # Initialize COB data collection like clean_dashboard does self.cob_data_received = 0 self.latest_cob_data = {} # Store all COB snapshots for heatmap generation self.cob_snapshots = deque() self.price_data = [] # For price line chart self.start_time = None self.subscriber_id = None def _tick_callback(self, tick: MarketTick): """Callback function to receive ticks from the DataProvider.""" if self.start_time is None: self.start_time = datetime.now() logger.info(f"Started collecting ticks at {self.start_time}") # Store all ticks self.ticks.append(tick) def _cob_data_callback(self, symbol: str, cob_data: dict): """Callback function to receive COB data from the DataProvider.""" self.cob_data_received += 1 self.latest_cob_data[symbol] = cob_data # Store the complete COB snapshot for heatmap generation if 'bids' in cob_data and 'asks' in cob_data: snapshot = { 'timestamp': cob_data.get('timestamp', datetime.now()), 'bids': cob_data['bids'], 'asks': cob_data['asks'], 'stats': cob_data.get('stats', {}) } self.cob_snapshots.append(snapshot) # Convert COB data to tick-like format for analysis if 'stats' in cob_data and 'mid_price' in cob_data['stats']: mid_price = cob_data['stats']['mid_price'] if mid_price > 0: # Store price data for line chart self.price_data.append({ 'timestamp': cob_data.get('timestamp', datetime.now()), 'price': mid_price }) # Create a synthetic tick from COB data synthetic_tick = MarketTick( symbol=symbol, timestamp=cob_data.get('timestamp', datetime.now()), price=mid_price, volume=cob_data.get('stats', {}).get('total_volume', 0), quantity=0, # Not available in COB data side='unknown', # COB data doesn't have side info trade_id=f"cob_{self.cob_data_received}", is_buyer_maker=False, raw_data=cob_data ) self.ticks.append(synthetic_tick) if self.cob_data_received % 10 == 0: # Log every 10th update logger.info(f"COB update #{self.cob_data_received}: {symbol} @ ${mid_price:.2f}") async def run_test(self): """Run the data collection and plotting test.""" logger.info(f"Starting COB stability test for {self.symbol} for {self.duration.total_seconds()} seconds...") # Initialize COB collection like clean_dashboard does try: logger.info("Starting COB collection in data provider...") self.data_provider.start_cob_collection() logger.info("Started COB collection in data provider") # Subscribe to COB updates logger.info("Subscribing to COB data updates...") self.data_provider.subscribe_to_cob(self._cob_data_callback) logger.info("Subscribed to COB data updates from data provider") except Exception as e: logger.error(f"Failed to start COB collection or subscribe: {e}") # Subscribe to ticks as fallback try: self.subscriber_id = self.data_provider.subscribe_to_ticks(self._tick_callback, symbols=[self.symbol]) logger.info("Subscribed to tick data as fallback") except Exception as e: logger.warning(f"Failed to subscribe to ticks: {e}") # Start the data provider's real-time streaming try: await self.data_provider.start_real_time_streaming() logger.info("Started real-time streaming") except Exception as e: logger.error(f"Failed to start real-time streaming: {e}") # Collect data for the specified duration self.start_time = datetime.now() while datetime.now() - self.start_time < self.duration: await asyncio.sleep(1) logger.info(f"Collected {len(self.ticks)} ticks so far...") # Stop streaming and unsubscribe await self.data_provider.stop_real_time_streaming() self.data_provider.unsubscribe_from_ticks(self.subscriber_id) logger.info(f"Finished collecting data. Total ticks: {len(self.ticks)}") # Plot the results if self.price_data and self.cob_snapshots: self.create_price_heatmap_chart() elif self.ticks: self._create_simple_price_chart() else: logger.warning("No data was collected. Cannot generate plot.") def create_price_heatmap_chart(self): """Create a visualization with price chart and order book heatmap.""" if not self.price_data or not self.cob_snapshots: logger.warning("Insufficient data to plot.") return logger.info(f"Creating price and order book heatmap chart...") # Prepare data price_df = pd.DataFrame(self.price_data) price_df['timestamp'] = pd.to_datetime(price_df['timestamp']) # Extract order book data for heatmap heatmap_data = [] for snapshot in self.cob_snapshots: timestamp = snapshot['timestamp'] for side in ['bids', 'asks']: for order in snapshot[side]: bucketed_price = round(order['price'] / self.price_granularity) * self.price_granularity heatmap_data.append({ 'time': timestamp, 'price': bucketed_price, 'size': order['size'], 'side': side }) heatmap_df = pd.DataFrame(heatmap_data) # Create plot fig, ax1 = plt.subplots(figsize=(16, 8)) # Plot price line ax1.plot(price_df['timestamp'], price_df['price'], 'cyan', linewidth=1, label='Price') # Prepare heatmap for side, cmap in zip(['bids', 'asks'], ['Greens', 'Reds']): side_df = heatmap_df[heatmap_df['side'] == side] if not side_df.empty: hist, xedges, yedges = np.histogram2d( side_df['time'].astype(np.int64) // 10**9, side_df['price'], bins=[np.unique(side_df['time'].astype(np.int64) // 10**9), np.arange(price_df['price'].min(), price_df['price'].max(), self.price_granularity)], weights=side_df['size'] ) ax1.pcolormesh(pd.to_datetime(xedges, unit='s'), yedges, hist.T, cmap=cmap, alpha=0.5) # Enhance plot ax1.set_title(f'Price Chart with Order Book Heatmap - {self.symbol}') ax1.set_xlabel('Time') ax1.set_ylabel('Price (USDT)') ax1.legend(loc='upper left') ax1.grid(True, alpha=0.3) plt.tight_layout() plot_filename = f"price_heatmap_chart_{self.symbol.replace('/', '_')}_{datetime.now():%Y%m%d_%H%M%S}.png" plt.savefig(plot_filename, dpi=150) logger.info(f"Price and heatmap chart saved to {plot_filename}") plt.show() def _create_simple_price_chart(self): """Create a simple price chart as fallback""" logger.info("Creating simple price chart as fallback...") prices = [] times = [] for tick in self.ticks: if tick.price > 0: prices.append(tick.price) times.append(tick.timestamp) if not prices: logger.warning("No price data to plot") return fig, ax = plt.subplots(figsize=(15, 8)) ax.plot(pd.to_datetime(times), prices, 'cyan', linewidth=1) ax.set_title(f'Price Chart - {self.symbol}') ax.set_xlabel('Time') ax.set_ylabel('Price (USDT)') fig.autofmt_xdate() plot_filename = f"cob_price_chart_{self.symbol.replace('/', '_')}_{datetime.now():%Y%m%d_%H%M%S}.png" plt.savefig(plot_filename) logger.info(f"Price chart saved to {plot_filename}") plt.show() async def main(symbol='ETHUSDT', duration_seconds=15): """Main function to run the COB test with configurable parameters. Args: symbol: Trading symbol (default: ETHUSDT) duration_seconds: Test duration in seconds (default: 15) """ logger.info(f"Starting COB test with symbol={symbol}, duration={duration_seconds}s") tester = COBStabilityTester(symbol=symbol, duration_seconds=duration_seconds) await tester.run_test() if __name__ == "__main__": import sys # Parse command line arguments symbol = 'ETHUSDT' # Default duration = 15 # Default if len(sys.argv) > 1: symbol = sys.argv[1] if len(sys.argv) > 2: try: duration = int(sys.argv[2]) except ValueError: logger.warning(f"Invalid duration '{sys.argv[2]}', using default 15 seconds") logger.info(f"Configuration: Symbol={symbol}, Duration={duration}s") logger.info(f"Granularity: {'1 USD for ETH' if 'ETH' in symbol.upper() else '10 USD for BTC' if 'BTC' in symbol.upper() else '1 USD default'}") try: asyncio.run(main(symbol, duration)) except KeyboardInterrupt: logger.info("Test interrupted by user.")