467 lines
18 KiB
Python
467 lines
18 KiB
Python
"""
|
|
Drone Formation Flight Simulation Server
|
|
Real-time 6-DOF physics and 5-Drone Swarm Formation calculation backend with WebSockets & Starlette.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import math
|
|
import time
|
|
import os
|
|
import contextlib
|
|
import sys
|
|
from typing import Dict, Any, List
|
|
import uvicorn
|
|
from starlette.applications import Starlette
|
|
from starlette.routing import Route, WebSocketRoute, Mount
|
|
from starlette.staticfiles import StaticFiles
|
|
from starlette.responses import FileResponse, JSONResponse
|
|
from starlette.websockets import WebSocket, WebSocketDisconnect
|
|
|
|
# Directory setup
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
STATIC_DIR = os.path.join(BASE_DIR, "static")
|
|
|
|
os.makedirs(STATIC_DIR, exist_ok=True)
|
|
os.makedirs(os.path.join(STATIC_DIR, "css"), exist_ok=True)
|
|
os.makedirs(os.path.join(STATIC_DIR, "js"), exist_ok=True)
|
|
|
|
|
|
class Vector3:
|
|
def __init__(self, x=0.0, y=0.0, z=0.0):
|
|
self.x = float(x)
|
|
self.y = float(y)
|
|
self.z = float(z)
|
|
|
|
def to_dict(self):
|
|
return {"x": round(self.x, 3), "y": round(self.y, 3), "z": round(self.z, 3)}
|
|
|
|
def dist_to(self, other: 'Vector3') -> float:
|
|
return math.sqrt((self.x - other.x)**2 + (self.y - other.y)**2 + (self.z - other.z)**2)
|
|
|
|
|
|
class DroneState:
|
|
def __init__(self, drone_id: str, name: str, callsign: str, initial_pos: Vector3, color: str, is_leader: bool = False):
|
|
self.id = drone_id
|
|
self.name = name
|
|
self.callsign = callsign
|
|
self.pos = initial_pos
|
|
self.target_pos = Vector3(initial_pos.x, initial_pos.y, initial_pos.z)
|
|
self.vel = Vector3(0.0, 0.0, 0.0)
|
|
self.pitch = 0.0 # radians
|
|
self.roll = 0.0 # radians
|
|
self.yaw = 0.0 # radians
|
|
self.rotor_rpm = [5200.0, 5200.0, 5200.0, 5200.0]
|
|
self.battery = 98.5 # percentage
|
|
self.armed = True
|
|
self.color = color
|
|
self.is_leader = is_leader
|
|
self.status = "AIRBORNE - FORMATION OK"
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
"id": self.id,
|
|
"name": self.name,
|
|
"callsign": self.callsign,
|
|
"position": self.pos.to_dict(),
|
|
"target": self.target_pos.to_dict(),
|
|
"velocity": self.vel.to_dict(),
|
|
"speed": round(math.sqrt(self.vel.x**2 + self.vel.y**2 + self.vel.z**2) * 3.6, 1), # km/h
|
|
"rotation": {
|
|
"pitch": round(math.degrees(self.pitch), 1),
|
|
"roll": round(math.degrees(self.roll), 1),
|
|
"yaw": round(math.degrees(self.yaw), 1)
|
|
},
|
|
"rotors": [round(r, 0) for r in self.rotor_rpm],
|
|
"battery": round(self.battery, 1),
|
|
"armed": self.armed,
|
|
"color": self.color,
|
|
"is_leader": self.is_leader,
|
|
"status": self.status
|
|
}
|
|
|
|
|
|
class FormationEngine:
|
|
"""Manages 5-Drone swarm flight physics and dynamic formation geometry"""
|
|
def __init__(self):
|
|
# 5 Drones Configuration
|
|
self.drones: List[DroneState] = [
|
|
DroneState("alpha_1", "Leader Drone", "VIPER-01", Vector3(0.0, 8.0, 0.0), "#00f0ff", is_leader=True),
|
|
DroneState("alpha_2", "Wingman 1 (R1)", "GHOST-02", Vector3(5.0, 8.0, -5.0), "#ff0055"),
|
|
DroneState("alpha_3", "Wingman 2 (L1)", "PHANTOM-03", Vector3(-5.0, 8.0, -5.0), "#00ff88"),
|
|
DroneState("alpha_4", "Wingman 3 (R2)", "SPECTER-04", Vector3(10.0, 8.0, -10.0),"#ffb700"),
|
|
DroneState("alpha_5", "Wingman 4 (L2)", "SHADOW-05", Vector3(-10.0, 8.0,-10.0),"#b026ff")
|
|
]
|
|
|
|
self.leader = self.drones[0]
|
|
self.formation_mode = "V_FORMATION" # V_FORMATION, SIDE_BY_SIDE, TANDEM, DIAMOND, ORBIT, VERTICAL, ECHELON
|
|
self.formation_distance = 6.0
|
|
self.formation_height_offset = 0.0
|
|
self.speed_multiplier = 1.0
|
|
self.autonomous_mode = False
|
|
self.auto_pattern = "FIGURE_EIGHT" # FIGURE_EIGHT, CIRCLE, WAYPOINT_PATROL
|
|
|
|
self.orbit_angle = 0.0
|
|
self.sim_time = 0.0
|
|
self.clients: List[WebSocket] = []
|
|
|
|
def set_target(self, x: float, y: float, z: float):
|
|
"""Set manual flight target for the leader drone"""
|
|
self.leader.target_pos.x = max(-80.0, min(80.0, float(x)))
|
|
self.leader.target_pos.y = max(1.5, min(40.0, float(y)))
|
|
self.leader.target_pos.z = max(-80.0, min(80.0, float(z)))
|
|
|
|
def calculate_follower_targets(self) -> List[Vector3]:
|
|
"""Calculates target 3D coordinates for all 4 wingmen based on formation geometry and leader state"""
|
|
lx, ly, lz = self.leader.pos.x, self.leader.pos.y, self.leader.pos.z
|
|
yaw = self.leader.yaw
|
|
dist = self.formation_distance
|
|
h_off = self.formation_height_offset
|
|
|
|
# Unit directional vectors based on leader yaw
|
|
# Forward vector
|
|
fx = math.sin(yaw)
|
|
fz = math.cos(yaw)
|
|
# Right vector
|
|
rx = math.cos(yaw)
|
|
rz = -math.sin(yaw)
|
|
|
|
targets = []
|
|
|
|
if self.formation_mode == "V_FORMATION":
|
|
# Classic 5-Drone V-Shape
|
|
offsets = [
|
|
# (Right factor, Forward factor, Height factor)
|
|
(1.0, -1.0, 0.0), # Drone 2: Starboard Inner
|
|
(-1.0, -1.0, 0.0), # Drone 3: Port Inner
|
|
(2.0, -2.0, 0.5), # Drone 4: Starboard Outer
|
|
(-2.0, -2.0, 0.5), # Drone 5: Port Outer
|
|
]
|
|
for r_fac, f_fac, h_fac in offsets:
|
|
tx = lx + (rx * r_fac + fx * f_fac) * dist
|
|
tz = lz + (rz * r_fac + fz * f_fac) * dist
|
|
ty = ly + h_off + (h_fac * h_off)
|
|
targets.append(Vector3(tx, ty, tz))
|
|
|
|
elif self.formation_mode == "SIDE_BY_SIDE":
|
|
# 5-Drone Wide Line Abreast: L2, L1, Leader, R1, R2
|
|
offsets = [
|
|
(1.0, 0.0), # D2: Right 1
|
|
(-1.0, 0.0), # D3: Left 1
|
|
(2.0, 0.0), # D4: Right 2
|
|
(-2.0, 0.0), # D5: Left 2
|
|
]
|
|
for r_fac, f_fac in offsets:
|
|
tx = lx + (rx * r_fac) * dist
|
|
tz = lz + (rz * r_fac) * dist
|
|
targets.append(Vector3(tx, ly + h_off, tz))
|
|
|
|
elif self.formation_mode == "TANDEM":
|
|
# 5-Drone Trail / Single Column In-Line behind Leader
|
|
offsets = [-1.0, -2.0, -3.0, -4.0]
|
|
for i, f_fac in enumerate(offsets):
|
|
tx = lx + (fx * f_fac) * dist
|
|
tz = lz + (fz * f_fac) * dist
|
|
ty = ly + (h_off * (i + 1) * 0.3)
|
|
targets.append(Vector3(tx, ty, tz))
|
|
|
|
elif self.formation_mode == "DIAMOND":
|
|
# 5-Drone Diamond / Arrowhead formation
|
|
offsets = [
|
|
(1.2, -0.8, 0.0), # D2: Right Wing
|
|
(-1.2, -0.8, 0.0), # D3: Left Wing
|
|
(0.0, -1.6, 0.0), # D4: Center Rear
|
|
(0.0, -2.6, 0.5), # D5: Tail High-Guard
|
|
]
|
|
for r_fac, f_fac, h_fac in offsets:
|
|
tx = lx + (rx * r_fac + fx * f_fac) * dist
|
|
tz = lz + (rz * r_fac + fz * f_fac) * dist
|
|
ty = ly + h_off + (h_fac * 1.5)
|
|
targets.append(Vector3(tx, ty, tz))
|
|
|
|
elif self.formation_mode == "ORBIT":
|
|
# 4 Wingmen dynamically orbiting around Leader at 90-degree spacing
|
|
for i in range(4):
|
|
angle = self.orbit_angle + (i * math.pi * 0.5)
|
|
ox = math.cos(angle) * dist
|
|
oz = math.sin(angle) * dist
|
|
oy = math.sin(angle * 2.0) * 1.5
|
|
targets.append(Vector3(lx + ox, ly + h_off + oy, lz + oz))
|
|
|
|
elif self.formation_mode == "ECHELON_RIGHT":
|
|
# Echelon Right (45 deg diagonal trail to the right)
|
|
for i in range(1, 5):
|
|
tx = lx + (rx * i - fx * i) * (dist * 0.75)
|
|
tz = lz + (rz * i - fz * i) * (dist * 0.75)
|
|
targets.append(Vector3(tx, ly + h_off + (i * 0.4), tz))
|
|
|
|
elif self.formation_mode == "VERTICAL":
|
|
# 3D Stack / Ladder
|
|
offsets = [
|
|
(0.0, 0.0, 1.0), # D2: Top +1
|
|
(0.0, 0.0, -1.0), # D3: Bottom -1
|
|
(0.5, -0.5, 2.0), # D4: High Rear
|
|
(-0.5, -0.5, -2.0) # D5: Low Rear
|
|
]
|
|
for r_fac, f_fac, h_fac in offsets:
|
|
tx = lx + (rx * r_fac + fx * f_fac) * (dist * 0.6)
|
|
tz = lz + (rz * r_fac + fz * f_fac) * (dist * 0.6)
|
|
ty = ly + (h_fac * dist * 0.6) + h_off
|
|
targets.append(Vector3(tx, ty, tz))
|
|
|
|
else:
|
|
# Fallback
|
|
for i in range(1, 5):
|
|
targets.append(Vector3(lx + i * 4.0, ly, lz))
|
|
|
|
return targets
|
|
|
|
def update_physics(self, dt: float):
|
|
"""Simulate realistic 6-DOF drone dynamics for all 5 drones"""
|
|
self.sim_time += dt
|
|
self.orbit_angle += dt * 1.2
|
|
|
|
# Autonomous flight patterns if enabled
|
|
if self.autonomous_mode:
|
|
t = self.sim_time * 0.35 * self.speed_multiplier
|
|
if self.auto_pattern == "FIGURE_EIGHT":
|
|
self.leader.target_pos.x = math.sin(t) * 35.0
|
|
self.leader.target_pos.z = math.sin(t * 2.0) * 20.0
|
|
self.leader.target_pos.y = 8.0 + math.sin(t * 3.0) * 3.0
|
|
elif self.auto_pattern == "CIRCLE":
|
|
self.leader.target_pos.x = math.cos(t) * 30.0
|
|
self.leader.target_pos.z = math.sin(t) * 30.0
|
|
self.leader.target_pos.y = 10.0 + math.cos(t * 2.0) * 2.5
|
|
elif self.auto_pattern == "WAYPOINT_PATROL":
|
|
wps = [Vector3(-25, 8, -25), Vector3(25, 12, -25), Vector3(25, 8, 25), Vector3(-25, 14, 25)]
|
|
idx = int((t * 0.5) % 4)
|
|
next_idx = (idx + 1) % 4
|
|
blend = (t * 0.5) % 1.0
|
|
p1 = wps[idx]
|
|
p2 = wps[next_idx]
|
|
self.leader.target_pos.x = p1.x + (p2.x - p1.x) * blend
|
|
self.leader.target_pos.y = p1.y + (p2.y - p1.y) * blend
|
|
self.leader.target_pos.z = p1.z + (p2.z - p1.z) * blend
|
|
|
|
# 1. Simulate Leader Drone (Moves toward target)
|
|
self._simulate_drone(self.leader, self.leader.target_pos, dt, responsiveness=3.2 * self.speed_multiplier)
|
|
|
|
# 2. Simulate 4 Follower Drones
|
|
follower_targets = self.calculate_follower_targets()
|
|
for idx, target in enumerate(follower_targets):
|
|
drone = self.drones[idx + 1]
|
|
drone.target_pos = target
|
|
self._simulate_drone(drone, target, dt, responsiveness=3.6 * self.speed_multiplier)
|
|
|
|
# Sync status
|
|
err = drone.pos.dist_to(target)
|
|
if err < 0.6:
|
|
drone.status = "FORMATION LOCKED"
|
|
elif err < 2.5:
|
|
drone.status = "TRACKING"
|
|
else:
|
|
drone.status = "REJOINING"
|
|
|
|
# Battery consumption
|
|
drone.battery = max(5.0, drone.battery - dt * 0.012)
|
|
|
|
self.leader.battery = max(5.0, self.leader.battery - dt * 0.012)
|
|
|
|
def _simulate_drone(self, drone: DroneState, target: Vector3, dt: float, responsiveness: float):
|
|
"""Critically damped spring-damper controller for drone navigation + attitude banking"""
|
|
ex = target.x - drone.pos.x
|
|
ey = target.y - drone.pos.y
|
|
ez = target.z - drone.pos.z
|
|
|
|
k_p = responsiveness * 3.5
|
|
k_d = 2.4 * math.sqrt(k_p)
|
|
|
|
acc_x = (ex * k_p) - (drone.vel.x * k_d)
|
|
acc_y = (ey * (k_p * 1.2)) - (drone.vel.y * k_d)
|
|
acc_z = (ez * k_p) - (drone.vel.z * k_d)
|
|
|
|
max_acc = 22.0
|
|
acc_mag = math.sqrt(acc_x**2 + acc_y**2 + acc_z**2)
|
|
if acc_mag > max_acc:
|
|
acc_x = (acc_x / acc_mag) * max_acc
|
|
acc_y = (acc_y / acc_mag) * max_acc
|
|
acc_z = (acc_z / acc_mag) * max_acc
|
|
|
|
drone.vel.x += acc_x * dt
|
|
drone.vel.y += acc_y * dt
|
|
drone.vel.z += acc_z * dt
|
|
|
|
drone.pos.x += drone.vel.x * dt
|
|
drone.pos.y += drone.vel.y * dt
|
|
drone.pos.z += drone.vel.z * dt
|
|
|
|
# Yaw alignment towards velocity
|
|
horizontal_speed = math.sqrt(drone.vel.x**2 + drone.vel.z**2)
|
|
if horizontal_speed > 0.3:
|
|
desired_yaw = math.atan2(drone.vel.x, drone.vel.z)
|
|
diff = (desired_yaw - drone.yaw + math.pi) % (2 * math.pi) - math.pi
|
|
drone.yaw += diff * min(1.0, dt * 4.0)
|
|
|
|
# Body-frame pitch & roll
|
|
cos_y = math.cos(drone.yaw)
|
|
sin_y = math.sin(drone.yaw)
|
|
forward_acc = acc_z * cos_y + acc_x * sin_y
|
|
lateral_acc = acc_x * cos_y - acc_z * sin_y
|
|
|
|
desired_pitch = -math.atan2(forward_acc, 9.81) * 0.7
|
|
desired_roll = math.atan2(lateral_acc, 9.81) * 0.7
|
|
|
|
max_tilt = math.radians(35.0)
|
|
desired_pitch = max(-max_tilt, min(max_tilt, desired_pitch))
|
|
desired_roll = max(-max_tilt, min(max_tilt, desired_roll))
|
|
|
|
drone.pitch += (desired_pitch - drone.pitch) * min(1.0, dt * 7.0)
|
|
drone.roll += (desired_roll - drone.roll) * min(1.0, dt * 7.0)
|
|
|
|
base_rpm = 5000.0 + (acc_y * 120.0)
|
|
tilt_factor = math.sqrt(drone.pitch**2 + drone.roll**2) * 1500.0
|
|
drone.rotor_rpm[0] = max(1000.0, base_rpm + tilt_factor + drone.roll * 800.0)
|
|
drone.rotor_rpm[1] = max(1000.0, base_rpm + tilt_factor - drone.roll * 800.0)
|
|
drone.rotor_rpm[2] = max(1000.0, base_rpm + tilt_factor + drone.pitch * 800.0)
|
|
drone.rotor_rpm[3] = max(1000.0, base_rpm + tilt_factor - drone.pitch * 800.0)
|
|
|
|
def get_state_payload(self) -> Dict[str, Any]:
|
|
return {
|
|
"timestamp": time.time(),
|
|
"count": len(self.drones),
|
|
"leader": self.leader.to_dict(),
|
|
"drones": [d.to_dict() for d in self.drones],
|
|
"formation": {
|
|
"mode": self.formation_mode,
|
|
"distance": round(self.formation_distance, 2),
|
|
"target_distance": self.formation_distance,
|
|
"height_offset": self.formation_height_offset,
|
|
"autonomous": self.autonomous_mode,
|
|
"auto_pattern": self.auto_pattern,
|
|
"speed_multiplier": round(self.speed_multiplier, 2)
|
|
}
|
|
}
|
|
|
|
|
|
engine = FormationEngine()
|
|
|
|
|
|
# Starlette App Endpoints
|
|
async def homepage(request):
|
|
index_file = os.path.join(STATIC_DIR, "index.html")
|
|
return FileResponse(index_file)
|
|
|
|
|
|
async def api_status(request):
|
|
return JSONResponse(engine.get_state_payload())
|
|
|
|
|
|
async def websocket_endpoint(websocket: WebSocket):
|
|
await websocket.accept()
|
|
engine.clients.append(websocket)
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_text()
|
|
try:
|
|
msg = json.loads(data)
|
|
cmd_type = msg.get("type")
|
|
|
|
if cmd_type == "SET_TARGET":
|
|
x = msg.get("x", 0.0)
|
|
y = msg.get("y", 8.0)
|
|
z = msg.get("z", 0.0)
|
|
engine.set_target(x, y, z)
|
|
|
|
elif cmd_type == "SET_FORMATION":
|
|
if "mode" in msg:
|
|
engine.formation_mode = msg["mode"]
|
|
if "distance" in msg:
|
|
engine.formation_distance = float(msg["distance"])
|
|
if "height_offset" in msg:
|
|
engine.formation_height_offset = float(msg["height_offset"])
|
|
|
|
elif cmd_type == "SET_AUTONOMOUS":
|
|
engine.autonomous_mode = bool(msg.get("enabled", False))
|
|
if "pattern" in msg:
|
|
engine.auto_pattern = msg["pattern"]
|
|
|
|
elif cmd_type == "SET_SPEED":
|
|
engine.speed_multiplier = float(msg.get("speed", 1.0))
|
|
|
|
elif cmd_type == "RESET_POSITION":
|
|
engine.leader.pos = Vector3(0.0, 8.0, 0.0)
|
|
engine.leader.target_pos = Vector3(0.0, 8.0, 0.0)
|
|
engine.leader.vel = Vector3(0.0, 0.0, 0.0)
|
|
# Reset wingmen around leader
|
|
for idx, d in enumerate(engine.drones[1:], 1):
|
|
d.pos = Vector3(idx * 4.0, 8.0, -idx * 3.0)
|
|
d.vel = Vector3(0.0, 0.0, 0.0)
|
|
|
|
except json.JSONDecodeError:
|
|
pass
|
|
except WebSocketDisconnect:
|
|
if websocket in engine.clients:
|
|
engine.clients.remove(websocket)
|
|
except Exception as e:
|
|
if websocket in engine.clients:
|
|
engine.clients.remove(websocket)
|
|
|
|
|
|
async def physics_loop():
|
|
"""High-frequency background simulation loop (60 FPS)"""
|
|
last_time = time.time()
|
|
while True:
|
|
now = time.time()
|
|
dt = min(0.05, now - last_time)
|
|
last_time = now
|
|
|
|
engine.update_physics(dt)
|
|
|
|
if engine.clients:
|
|
payload = json.dumps(engine.get_state_payload())
|
|
dead_clients = []
|
|
for client in engine.clients:
|
|
try:
|
|
await client.send_text(payload)
|
|
except Exception:
|
|
dead_clients.append(client)
|
|
for dead in dead_clients:
|
|
if dead in engine.clients:
|
|
engine.clients.remove(dead)
|
|
|
|
await asyncio.sleep(1.0 / 60.0)
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def lifespan(app):
|
|
sim_task = asyncio.create_task(physics_loop())
|
|
yield
|
|
sim_task.cancel()
|
|
try:
|
|
await sim_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
|
|
routes = [
|
|
Route("/", homepage),
|
|
WebSocketRoute("/ws", websocket_endpoint),
|
|
Route("/api/status", api_status),
|
|
Mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
|
|
]
|
|
|
|
app = Starlette(debug=True, routes=routes, lifespan=lifespan)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if hasattr(sys.stdout, "reconfigure"):
|
|
try:
|
|
sys.stdout.reconfigure(encoding="utf-8")
|
|
except Exception:
|
|
pass
|
|
print("================================================================")
|
|
print(" [5-DRONE FORMATION FLIGHT SIMULATOR] Three.js + Python")
|
|
print(" Server starting at: http://127.0.0.1:8000")
|
|
print("================================================================")
|
|
uvicorn.run("app:app", host="127.0.0.1", port=8000, log_level="info", reload=False)
|