feat(simulator): virtual drone fleet data generator
One process per drone: seeded patrol kinematics, sensor streams at realistic rates, detection events, hot current/ blocks sealed into Hive-partitioned Parquet, and the compact UDP state broadcast with peer RSSI derivation. Scales via docker compose --scale.
This commit is contained in:
@@ -0,0 +1,8 @@
|
||||
"""Virtual drone: a data generator that exercises the exact on-board pipeline.
|
||||
|
||||
One process simulates one drone: kinematics along a patrol route, sensor
|
||||
streams at realistic rates, detection events, the Hive-partitioned Parquet
|
||||
write path (current/ blocks -> sealed files), and the UDP state broadcast.
|
||||
"""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
@@ -0,0 +1,68 @@
|
||||
"""State broadcast over UDP: the compact pose frame from the sync design.
|
||||
|
||||
Frame layout (little-endian, 46 bytes):
|
||||
|
||||
magic 2s b"SH"
|
||||
version B
|
||||
drone_id 8s zero-padded ascii
|
||||
ts_ns q epoch nanoseconds
|
||||
pos_mm 3i position, millimeters (quantized on the wire only)
|
||||
att_cdeg 3h roll/pitch/yaw, centi-degrees
|
||||
vel_cms 3h velocity, cm/s
|
||||
frame_ref B
|
||||
flags B
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import socket
|
||||
import struct
|
||||
from typing import Any
|
||||
|
||||
from .flight import Pose
|
||||
|
||||
FRAME = struct.Struct("<2sB8sq3i3h3hBB")
|
||||
MAGIC = b"SH"
|
||||
VERSION = 1
|
||||
|
||||
|
||||
def encode_state(drone_id: str, pose: Pose, flags: int = 0) -> bytes:
|
||||
return FRAME.pack(
|
||||
MAGIC,
|
||||
VERSION,
|
||||
drone_id.encode()[:8].ljust(8, b"\0"),
|
||||
pose.ts_ns,
|
||||
int(pose.x * 1000), int(pose.y * 1000), int(pose.z * 1000),
|
||||
int(pose.roll * 100), int(pose.pitch * 100), int(pose.yaw * 100),
|
||||
int(pose.vx * 100), int(pose.vy * 100), int(pose.vz * 100),
|
||||
0,
|
||||
flags,
|
||||
)
|
||||
|
||||
|
||||
def decode_state(payload: bytes) -> dict[str, Any] | None:
|
||||
if len(payload) != FRAME.size:
|
||||
return None
|
||||
magic, version, drone_id, ts_ns, px, py, pz, r, p, y, vx, vy, vz, frame_ref, flags = FRAME.unpack(payload)
|
||||
if magic != MAGIC or version != VERSION:
|
||||
return None
|
||||
return {
|
||||
"ts_ns": ts_ns,
|
||||
"peer_id": drone_id.rstrip(b"\0").decode(),
|
||||
"pos_x": px / 1000.0, "pos_y": py / 1000.0, "pos_z": pz / 1000.0,
|
||||
"roll": r / 100.0, "pitch": p / 100.0, "yaw": y / 100.0,
|
||||
"vel_x": vx / 100.0, "vel_y": vy / 100.0, "vel_z": vz / 100.0,
|
||||
"frame_ref": frame_ref,
|
||||
"flags": flags,
|
||||
}
|
||||
|
||||
|
||||
def open_socket(port: int) -> socket.socket:
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
|
||||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
if hasattr(socket, "SO_REUSEPORT"):
|
||||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
|
||||
sock.bind(("", port))
|
||||
sock.setblocking(False)
|
||||
return sock
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Runtime configuration, resolved once from the environment."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import socket
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
def _default_drone_id() -> str:
|
||||
"""Stable per-container id: works with `docker compose up --scale`."""
|
||||
return os.environ.get("DRONE_ID") or f"dr-{socket.gethostname()[-4:]}"
|
||||
|
||||
|
||||
def _default_flight_id() -> str:
|
||||
"""Short time-sortable partition alias; canonical UUIDv7 kept separately."""
|
||||
return os.environ.get("FLIGHT_ID") or time.strftime(
|
||||
"%Y%m%dT%H%MZ", time.gmtime()
|
||||
) + "-" + uuid.uuid4().hex[:4]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Config:
|
||||
drone_id: str = field(default_factory=_default_drone_id)
|
||||
flight_id: str = field(default_factory=_default_flight_id)
|
||||
flight_uuid: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
|
||||
data_dir: str = field(default_factory=lambda: os.environ.get("DATA_DIR", "./data"))
|
||||
seed: int = field(default_factory=lambda: int(os.environ.get("SEED", "42")))
|
||||
|
||||
# Simulation pacing
|
||||
duration_s: float = field(
|
||||
default_factory=lambda: float(os.environ.get("DURATION_S", "120"))
|
||||
)
|
||||
speedup: float = field(
|
||||
default_factory=lambda: float(os.environ.get("SPEEDUP", "1.0"))
|
||||
)
|
||||
|
||||
# Sensor rates (Hz, simulated time)
|
||||
imu_hz: float = field(default_factory=lambda: float(os.environ.get("IMU_HZ", "50")))
|
||||
baro_hz: float = 1.0
|
||||
temp_hz: float = 0.2
|
||||
battery_hz: float = 0.5
|
||||
detection_rate_hz: float = field(
|
||||
default_factory=lambda: float(os.environ.get("DETECTION_HZ", "0.05"))
|
||||
)
|
||||
|
||||
# State broadcast
|
||||
state_hz: float = field(
|
||||
default_factory=lambda: float(os.environ.get("STATE_HZ", "5"))
|
||||
)
|
||||
udp_port: int = field(default_factory=lambda: int(os.environ.get("UDP_PORT", "47000")))
|
||||
broadcast_addr: str = field(
|
||||
default_factory=lambda: os.environ.get("BROADCAST_ADDR", "255.255.255.255")
|
||||
)
|
||||
|
||||
# Patrol area (meters, local mission frame)
|
||||
area_size_m: float = field(
|
||||
default_factory=lambda: float(os.environ.get("AREA_SIZE_M", "400"))
|
||||
)
|
||||
cruise_alt_m: float = field(
|
||||
default_factory=lambda: float(os.environ.get("CRUISE_ALT_M", "60"))
|
||||
)
|
||||
cruise_speed_ms: float = field(
|
||||
default_factory=lambda: float(os.environ.get("CRUISE_SPEED_MS", "12"))
|
||||
)
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Kinematics: a seeded patrol route and the drone's motion along it."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import random
|
||||
from dataclasses import dataclass, replace
|
||||
|
||||
from .config import Config
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Pose:
|
||||
ts_ns: int
|
||||
x: float
|
||||
y: float
|
||||
z: float
|
||||
roll: float
|
||||
pitch: float
|
||||
yaw: float
|
||||
vx: float
|
||||
vy: float
|
||||
vz: float
|
||||
|
||||
|
||||
def make_waypoints(cfg: Config, rng: random.Random) -> list[tuple[float, float]]:
|
||||
"""Perimeter patrol with per-drone jitter so drones do not overlap."""
|
||||
s = cfg.area_size_m
|
||||
jitter = rng.uniform(0.05, 0.20) * s
|
||||
corners = [(jitter, jitter), (s - jitter, jitter), (s - jitter, s - jitter), (jitter, s - jitter)]
|
||||
start = rng.randrange(4)
|
||||
return corners[start:] + corners[:start]
|
||||
|
||||
|
||||
def initial_pose(cfg: Config, rng: random.Random, ts_ns: int) -> Pose:
|
||||
wps = make_waypoints(cfg, rng)
|
||||
x, y = wps[0]
|
||||
return Pose(ts_ns=ts_ns, x=x, y=y, z=cfg.cruise_alt_m, roll=0.0, pitch=0.0, yaw=0.0, vx=0.0, vy=0.0, vz=0.0)
|
||||
|
||||
|
||||
def step(pose: Pose, target: tuple[float, float], cfg: Config, rng: random.Random, dt: float, ts_ns: int) -> tuple[Pose, bool]:
|
||||
"""Advance one tick toward the target waypoint. Returns (pose, reached)."""
|
||||
dx, dy = target[0] - pose.x, target[1] - pose.y
|
||||
dist = math.hypot(dx, dy)
|
||||
reached = dist < cfg.cruise_speed_ms * dt * 2
|
||||
|
||||
speed = cfg.cruise_speed_ms * rng.uniform(0.9, 1.1)
|
||||
if dist > 1e-6:
|
||||
vx, vy = dx / dist * speed, dy / dist * speed
|
||||
else:
|
||||
vx, vy = 0.0, 0.0
|
||||
|
||||
yaw = math.degrees(math.atan2(vy, vx))
|
||||
# Small attitude noise around the commanded values
|
||||
roll = rng.gauss(0.0, 1.5)
|
||||
pitch = rng.gauss(2.0, 1.0)
|
||||
z = cfg.cruise_alt_m + rng.gauss(0.0, 0.5)
|
||||
vz = (z - pose.z) / dt if dt > 0 else 0.0
|
||||
|
||||
return (
|
||||
replace(
|
||||
pose,
|
||||
ts_ns=ts_ns,
|
||||
x=pose.x + vx * dt,
|
||||
y=pose.y + vy * dt,
|
||||
z=z,
|
||||
roll=roll,
|
||||
pitch=pitch,
|
||||
yaw=yaw,
|
||||
vx=vx,
|
||||
vy=vy,
|
||||
vz=vz,
|
||||
),
|
||||
reached,
|
||||
)
|
||||
@@ -0,0 +1,114 @@
|
||||
"""Virtual drone entry point: fly, sense, write, broadcast, listen, seal."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import select
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from .broadcast import decode_state, encode_state, open_socket
|
||||
from .config import Config
|
||||
from .flight import initial_pose, make_waypoints, step
|
||||
from .sensors import (
|
||||
baro_row,
|
||||
battery_row,
|
||||
detection_row,
|
||||
imu_row,
|
||||
state_row,
|
||||
temp_row,
|
||||
)
|
||||
from .writer import make_writers
|
||||
|
||||
TICK_HZ = 10.0
|
||||
SENSOR_NAMES = ["imu", "baro", "temp", "battery", "rssi"]
|
||||
|
||||
|
||||
def run(cfg: Config) -> None:
|
||||
rng = random.Random(f"{cfg.seed}:{cfg.drone_id}")
|
||||
root = Path(cfg.data_dir)
|
||||
writers = make_writers(root, cfg.flight_id, cfg.drone_id, SENSOR_NAMES)
|
||||
sock = open_socket(cfg.udp_port)
|
||||
|
||||
now_ns = time.time_ns()
|
||||
pose = initial_pose(cfg, rng, now_ns)
|
||||
waypoints = make_waypoints(cfg, rng)
|
||||
wp_index = 1
|
||||
|
||||
dt = 1.0 / TICK_HZ
|
||||
ticks_total = int(cfg.duration_s * TICK_HZ)
|
||||
# Emission accumulators: sensor name -> carry-over fraction of a sample
|
||||
acc = {name: 0.0 for name in ["imu", "baro", "temp", "battery", "state", "detection"]}
|
||||
rates = {
|
||||
"imu": cfg.imu_hz, "baro": cfg.baro_hz, "temp": cfg.temp_hz,
|
||||
"battery": cfg.battery_hz, "state": cfg.state_hz, "detection": cfg.detection_rate_hz,
|
||||
}
|
||||
peers_seen: dict[str, int] = {}
|
||||
frames_sent = 0
|
||||
|
||||
print(f"[{cfg.drone_id}] flight={cfg.flight_id} uuid={cfg.flight_uuid} "
|
||||
f"duration={cfg.duration_s}s speedup={cfg.speedup}x data={root}")
|
||||
|
||||
for tick in range(ticks_total):
|
||||
sim_elapsed = tick * dt
|
||||
ts_ns = now_ns + int(sim_elapsed * 1e9)
|
||||
|
||||
pose, reached = step(pose, waypoints[wp_index], cfg, rng, dt, ts_ns)
|
||||
if reached:
|
||||
wp_index = (wp_index + 1) % len(waypoints)
|
||||
|
||||
def due(name: str) -> int:
|
||||
acc[name] += rates[name] * dt
|
||||
n = int(acc[name])
|
||||
acc[name] -= n
|
||||
return n
|
||||
|
||||
for i in range(due("imu")):
|
||||
writers["imu"].append(imu_row(pose, rng, ts_ns + i * int(1e9 / max(cfg.imu_hz, 1))))
|
||||
for _ in range(due("baro")):
|
||||
writers["baro"].append(baro_row(pose, rng, ts_ns))
|
||||
for _ in range(due("temp")):
|
||||
writers["temp"].append(temp_row(pose, rng, ts_ns))
|
||||
for _ in range(due("battery")):
|
||||
writers["battery"].append(battery_row(pose, rng, ts_ns, sim_elapsed, cfg.duration_s))
|
||||
for _ in range(due("detection")):
|
||||
writers["detections"].append(detection_row(pose, rng, ts_ns))
|
||||
|
||||
for _ in range(due("state")):
|
||||
frame = encode_state(cfg.drone_id, pose)
|
||||
try:
|
||||
sock.sendto(frame, (cfg.broadcast_addr, cfg.udp_port))
|
||||
frames_sent += 1
|
||||
except OSError:
|
||||
pass # lossy link is part of the model
|
||||
writers["state"].append({**state_row(pose, ts_ns), "peer_id": cfg.drone_id, "direction": "sent"})
|
||||
|
||||
# Drain incoming peer frames; record them and derive RSSI-like signal
|
||||
while True:
|
||||
readable, _, _ = select.select([sock], [], [], 0)
|
||||
if not readable:
|
||||
break
|
||||
payload, _addr = sock.recvfrom(256)
|
||||
decoded = decode_state(payload)
|
||||
if decoded is None or decoded["peer_id"] == cfg.drone_id:
|
||||
continue
|
||||
peers_seen[decoded["peer_id"]] = decoded["ts_ns"]
|
||||
writers["state"].append({**decoded, "direction": "received"})
|
||||
dist = max(1.0, ((decoded["pos_x"] - pose.x) ** 2 + (decoded["pos_y"] - pose.y) ** 2) ** 0.5)
|
||||
writers["rssi"].append({
|
||||
"ts_ns": ts_ns,
|
||||
"peer_id": decoded["peer_id"],
|
||||
"rssi_dbm": -40.0 - 20.0 * (dist ** 0.5) / 10.0 + rng.gauss(0.0, 2.0),
|
||||
"distance_m": dist,
|
||||
})
|
||||
|
||||
time.sleep(dt / max(cfg.speedup, 0.01))
|
||||
|
||||
for writer in writers.values():
|
||||
writer.seal()
|
||||
print(f"[{cfg.drone_id}] done: {frames_sent} state frames sent, "
|
||||
f"{len(peers_seen)} peers seen {sorted(peers_seen)}; sealed to {root}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run(Config())
|
||||
@@ -0,0 +1,80 @@
|
||||
"""Sensor row generators. Pure functions: (pose, rng, ts) -> row dict.
|
||||
|
||||
Every row carries int64 epoch nanoseconds. Values are emitted at full float
|
||||
precision - quantization happens only in the broadcast payload, never in
|
||||
stored telemetry.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
from .flight import Pose
|
||||
|
||||
Row = dict[str, Any]
|
||||
|
||||
|
||||
def imu_row(pose: Pose, rng: random.Random, ts_ns: int) -> Row:
|
||||
return {
|
||||
"ts_ns": ts_ns,
|
||||
"accel_x": rng.gauss(0.0, 0.35),
|
||||
"accel_y": rng.gauss(0.0, 0.35),
|
||||
"accel_z": rng.gauss(-9.81, 0.25),
|
||||
"gyro_x": math.radians(rng.gauss(pose.roll, 0.8)),
|
||||
"gyro_y": math.radians(rng.gauss(pose.pitch, 0.8)),
|
||||
"gyro_z": math.radians(rng.gauss(0.0, 0.5)),
|
||||
}
|
||||
|
||||
|
||||
def baro_row(pose: Pose, rng: random.Random, ts_ns: int) -> Row:
|
||||
pressure = 101325.0 * math.exp(-pose.z / 8434.0) + rng.gauss(0.0, 4.0)
|
||||
return {"ts_ns": ts_ns, "pressure_pa": pressure, "alt_est_m": pose.z + rng.gauss(0.0, 0.3)}
|
||||
|
||||
|
||||
def temp_row(_pose: Pose, rng: random.Random, ts_ns: int) -> Row:
|
||||
return {
|
||||
"ts_ns": ts_ns,
|
||||
"cpu_c": 55.0 + rng.gauss(0.0, 3.0),
|
||||
"gpu_c": 62.0 + rng.gauss(0.0, 4.0),
|
||||
"ambient_c": 24.0 + rng.gauss(0.0, 0.5),
|
||||
}
|
||||
|
||||
|
||||
def battery_row(_pose: Pose, rng: random.Random, ts_ns: int, elapsed_s: float, duration_s: float) -> Row:
|
||||
level = max(0.0, 100.0 - 80.0 * elapsed_s / max(duration_s, 1.0) + rng.gauss(0.0, 0.2))
|
||||
return {"ts_ns": ts_ns, "level_pct": level, "voltage_v": 22.2 * (0.85 + 0.15 * level / 100.0)}
|
||||
|
||||
|
||||
DETECTION_CLASSES = ("vehicle", "person", "animal", "structure", "unknown")
|
||||
|
||||
|
||||
def detection_row(pose: Pose, rng: random.Random, ts_ns: int) -> Row:
|
||||
return {
|
||||
"ts_ns": ts_ns,
|
||||
"cls": rng.choice(DETECTION_CLASSES),
|
||||
"confidence": round(rng.uniform(0.42, 0.99), 3),
|
||||
"obj_x": pose.x + rng.uniform(-40.0, 40.0),
|
||||
"obj_y": pose.y + rng.uniform(-40.0, 40.0),
|
||||
"obj_z": rng.uniform(0.0, 5.0),
|
||||
"ego_x": pose.x,
|
||||
"ego_y": pose.y,
|
||||
"ego_z": pose.z,
|
||||
}
|
||||
|
||||
|
||||
def state_row(pose: Pose, ts_ns: int) -> Row:
|
||||
return {
|
||||
"ts_ns": ts_ns,
|
||||
"pos_x": pose.x,
|
||||
"pos_y": pose.y,
|
||||
"pos_z": pose.z,
|
||||
"roll": pose.roll,
|
||||
"pitch": pose.pitch,
|
||||
"yaw": pose.yaw,
|
||||
"vel_x": pose.vx,
|
||||
"vel_y": pose.vy,
|
||||
"vel_z": pose.vz,
|
||||
"frame_ref": 0,
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Hive-partitioned Parquet writer: the hot `current/` path plus sealing.
|
||||
|
||||
Layout (identical to the on-board and warehouse design):
|
||||
|
||||
dataset=<name>/flight=<id>/drone=<id>[/sensor=<type>]/
|
||||
year=YYYY/month=MM/day=DD/hour=HH/
|
||||
current/min_MM.parquet <- open window, small blocks
|
||||
data.parquet <- sealed on window close / shutdown
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
|
||||
Row = dict[str, Any]
|
||||
|
||||
|
||||
def _hour_dir(root: Path, dataset: str, flight: str, drone: str, sensor: str | None, ts_ns: int) -> Path:
|
||||
t = time.gmtime(ts_ns / 1e9)
|
||||
parts = [f"dataset={dataset}", f"flight={flight}", f"drone={drone}"]
|
||||
if sensor is not None:
|
||||
parts.append(f"sensor={sensor}")
|
||||
parts += [
|
||||
f"year={t.tm_year:04d}", f"month={t.tm_mon:02d}",
|
||||
f"day={t.tm_mday:02d}", f"hour={t.tm_hour:02d}",
|
||||
]
|
||||
return root.joinpath(*parts)
|
||||
|
||||
|
||||
class PartitionWriter:
|
||||
"""Buffers rows per minute and flushes small blocks into current/."""
|
||||
|
||||
def __init__(self, root: Path, dataset: str, flight: str, drone: str, sensor: str | None) -> None:
|
||||
self._root = root
|
||||
self._dataset = dataset
|
||||
self._flight = flight
|
||||
self._drone = drone
|
||||
self._sensor = sensor
|
||||
self._buffer: list[Row] = []
|
||||
self._minute: int | None = None
|
||||
self._touched: set[Path] = set()
|
||||
|
||||
def append(self, row: Row) -> None:
|
||||
minute = int(row["ts_ns"] // 60_000_000_000)
|
||||
if self._minute is not None and minute != self._minute:
|
||||
self._flush()
|
||||
self._minute = minute
|
||||
self._buffer.append(row)
|
||||
|
||||
def _flush(self) -> None:
|
||||
if not self._buffer or self._minute is None:
|
||||
return
|
||||
first_ts = self._buffer[0]["ts_ns"]
|
||||
hour_dir = _hour_dir(self._root, self._dataset, self._flight, self._drone, self._sensor, first_ts)
|
||||
current = hour_dir / "current"
|
||||
current.mkdir(parents=True, exist_ok=True)
|
||||
table = pa.Table.from_pylist(self._buffer)
|
||||
minute_of_hour = self._minute % 60
|
||||
pq.write_table(table, current / f"min_{minute_of_hour:02d}.parquet", compression="zstd")
|
||||
self._touched.add(hour_dir)
|
||||
self._buffer = []
|
||||
|
||||
def seal(self) -> None:
|
||||
"""Compact every touched hour's current/ blocks into one data.parquet."""
|
||||
self._flush()
|
||||
for hour_dir in sorted(self._touched):
|
||||
current = hour_dir / "current"
|
||||
blocks = sorted(current.glob("min_*.parquet"))
|
||||
if not blocks:
|
||||
continue
|
||||
table = pa.concat_tables([pq.read_table(b) for b in blocks])
|
||||
pq.write_table(table, hour_dir / "data.parquet", compression="zstd", compression_level=9)
|
||||
for b in blocks:
|
||||
b.unlink()
|
||||
current.rmdir()
|
||||
|
||||
|
||||
def make_writers(root: Path, flight: str, drone: str, sensor_names: list[str]) -> dict[str, PartitionWriter]:
|
||||
"""One telemetry writer per sensor + detections + state."""
|
||||
writers = {
|
||||
name: PartitionWriter(root, "telemetry", flight, drone, name)
|
||||
for name in sensor_names
|
||||
}
|
||||
writers["detections"] = PartitionWriter(root, "detections", flight, drone, None)
|
||||
writers["state"] = PartitionWriter(root, "state", flight, drone, None)
|
||||
return writers
|
||||
Reference in New Issue
Block a user