Spaces:
Sleeping
Sleeping
| from collections import defaultdict | |
| class OrderBook: | |
| def __init__(self, arrival_price=100.0): | |
| self.bids = defaultdict(list) | |
| self.asks = defaultdict(list) | |
| self.orders = {} | |
| self.arrival_price = arrival_price | |
| self.next_id = 0 | |
| self.total_traded_volume = 0.0 | |
| def add_order(self, side, price, size, agent_id): | |
| """ | |
| Add a limit order and immediately match it against the opposite side. | |
| Returns (order_id, fills) where: | |
| - order_id is the ID of the resting order (or None if fully filled) | |
| - fills is a list of fill dicts with keys: price, size, buyer, seller | |
| """ | |
| order_id = self.next_id | |
| self.next_id += 1 | |
| order = { | |
| "id": order_id, | |
| "side": side, | |
| "price": price, | |
| "size": size, | |
| "agent_id": agent_id, | |
| } | |
| self.orders[order_id] = order | |
| remaining = size | |
| fills = [] | |
| if side == "bid": | |
| while ( | |
| remaining > 0 and self.best_ask is not None and price >= self.best_ask | |
| ): | |
| current_best_ask = self.best_ask | |
| if current_best_ask not in self.asks or not self.asks[current_best_ask]: | |
| if current_best_ask in self.asks: | |
| del self.asks[current_best_ask] | |
| continue | |
| ask_order = self.asks[current_best_ask][0] | |
| fill = min(remaining, ask_order["size"]) | |
| fills.append( | |
| { | |
| "price": ask_order["price"], | |
| "size": fill, | |
| "buyer": agent_id, | |
| "seller": ask_order["agent_id"], | |
| } | |
| ) | |
| self.total_traded_volume += fill | |
| remaining -= fill | |
| ask_order["size"] -= fill | |
| if ask_order["size"] == 0: | |
| self.cancel_order(ask_order["id"]) | |
| if not self.asks[current_best_ask]: | |
| del self.asks[current_best_ask] | |
| if remaining > 0: | |
| order["size"] = remaining | |
| if price not in self.bids: | |
| self.bids[price] = [] | |
| self.bids[price].append(order) | |
| return order_id, fills | |
| else: | |
| if order_id in self.orders: | |
| del self.orders[order_id] | |
| return None, fills | |
| else: # side == 'ask' | |
| while ( | |
| remaining > 0 and self.best_bid is not None and price <= self.best_bid | |
| ): | |
| current_best_bid = self.best_bid | |
| if current_best_bid not in self.bids or not self.bids[current_best_bid]: | |
| if current_best_bid in self.bids: | |
| del self.bids[current_best_bid] | |
| continue | |
| bid_order = self.bids[current_best_bid][0] | |
| fill = min(remaining, bid_order["size"]) | |
| fills.append( | |
| { | |
| "price": bid_order["price"], | |
| "size": fill, | |
| "buyer": bid_order["agent_id"], | |
| "seller": agent_id, | |
| } | |
| ) | |
| self.total_traded_volume += fill | |
| remaining -= fill | |
| bid_order["size"] -= fill | |
| if bid_order["size"] == 0: | |
| self.cancel_order(bid_order["id"]) | |
| if not self.bids[current_best_bid]: | |
| del self.bids[current_best_bid] | |
| if remaining > 0: | |
| order["size"] = remaining | |
| if price not in self.asks: | |
| self.asks[price] = [] | |
| self.asks[price].append(order) | |
| return order_id, fills | |
| else: | |
| if order_id in self.orders: | |
| del self.orders[order_id] | |
| return None, fills | |
| def cancel_order(self, order_id): | |
| if order_id not in self.orders: | |
| return False | |
| order = self.orders.pop(order_id) | |
| level = self.bids if order["side"] == "bid" else self.asks | |
| level[order["price"]] = [ | |
| o for o in level[order["price"]] if o["id"] != order_id | |
| ] | |
| if not level[order["price"]]: | |
| del level[order["price"]] | |
| def best_bid(self): | |
| return max(self.bids.keys()) if self.bids else None | |
| def best_ask(self): | |
| return min(self.asks.keys()) if self.asks else None | |
| def mid(self): | |
| if self.best_bid and self.best_ask: | |
| return (self.best_bid + self.best_ask) / 2 | |
| return self.arrival_price | |
| def get_depth(self, levels=5): | |
| bid_levels = sorted(self.bids.keys(), reverse=True)[:levels] | |
| ask_levels = sorted(self.asks.keys())[:levels] | |
| return { | |
| "bid_prices": bid_levels, | |
| "bid_sizes": [sum(o["size"] for o in self.bids[p]) for p in bid_levels], | |
| "ask_prices": ask_levels, | |
| "ask_sizes": [sum(o["size"] for o in self.asks[p]) for p in ask_levels], | |
| } | |