drones-sim 0.2.0__py3-none-any.whl

drones_sim/sensors/imu.py

@@ -0,0 +1,196 @@
+"""IMU sensor simulator with realistic noise, bias, scale factor, and optional temperature effects.
+
+Consolidated from imu_ekf_simulation.py, imu_ekf_fusion_enhanced.py,
+imu_ekf_fusion_final.py, and imu_ekf_fusion_simplified.py.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy.spatial.transform import Rotation as R
+
+from ..state import VehicleState
+from ..trajectory import TrajectoryData
+from .models import SensorNoiseModel, TemperatureModel
+
+
+@dataclass
+class IMUConfig:
+    """Configuration for IMU sensor characteristics."""
+    accel_noise_std: float = 0.05       # m/s^2
+    accel_bias_range: float = 0.1       # m/s^2
+    accel_scale: tuple[float, float] = (0.98, 1.02)
+
+    gyro_noise_std: float = 0.01        # rad/s
+    gyro_bias_range: float = 0.005      # rad/s
+    gyro_scale: tuple[float, float] = (0.99, 1.01)
+
+    mag_noise_std: float = 0.5          # uT
+    mag_bias_range: float = 1.0         # uT
+    mag_scale: tuple[float, float] = (0.97, 1.03)
+
+    # Earth references
+    gravity: NDArray = field(default_factory=lambda: np.array([0.0, 0.0, 9.81]))
+    mag_field_ref: NDArray = field(default_factory=lambda: np.array([25.0, 5.0, -40.0]))
+
+    enable_temperature: bool = False
+
+
+@dataclass
+class IMUData:
+    """Container for simulated IMU readings."""
+    t: NDArray
+    accel: NDArray          # (N, 3)
+    gyro: NDArray           # (N, 3)
+    mag: NDArray            # (N, 3)
+    temperature: NDArray | None = None  # (N,) if temperature model enabled
+
+    # Ground-truth biases for validation
+    accel_bias: NDArray = field(default_factory=lambda: np.zeros(3))
+    gyro_bias: NDArray = field(default_factory=lambda: np.zeros(3))
+    mag_bias: NDArray = field(default_factory=lambda: np.zeros(3))
+
+
+@dataclass
+class IMUMeasurement:
+    """One timestamped 9-axis IMU measurement."""
+
+    acceleration: NDArray
+    angular_velocity: NDArray
+    magnetic_field: NDArray
+    temperature: float | None = None
+
+
+class IMUSimulator:
+    """Simulate 9-axis IMU readings from a known trajectory.
+
+    Supports:
+    - Additive white Gaussian noise
+    - Constant bias per axis
+    - Scale-factor errors
+    - Optional temperature-dependent drift
+    """
+
+    def __init__(self, config: IMUConfig | None = None, seed: int | None = 42):
+        self.cfg = config or IMUConfig()
+        self._seed = seed
+        seed_sequence = np.random.SeedSequence(seed)
+        accel_rng, gyro_rng, mag_rng, temp_rng = [
+            np.random.default_rng(child) for child in seed_sequence.spawn(4)
+        ]
+
+        self.accel_model = SensorNoiseModel(
+            noise_std=self.cfg.accel_noise_std,
+            bias_range=self.cfg.accel_bias_range,
+            scale_factor_range=self.cfg.accel_scale,
+            rng=accel_rng,
+        )
+        self.gyro_model = SensorNoiseModel(
+            noise_std=self.cfg.gyro_noise_std,
+            bias_range=self.cfg.gyro_bias_range,
+            scale_factor_range=self.cfg.gyro_scale,
+            rng=gyro_rng,
+        )
+        self.mag_model = SensorNoiseModel(
+            noise_std=self.cfg.mag_noise_std,
+            bias_range=self.cfg.mag_bias_range,
+            scale_factor_range=self.cfg.mag_scale,
+            rng=mag_rng,
+        )
+        self.temp_model = (
+            TemperatureModel(rng=temp_rng) if self.cfg.enable_temperature else None
+        )
+
+    def reset(self, seed: int | None = None) -> None:
+        """Reset biases and local random streams for a new episode."""
+        self.__init__(self.cfg, self._seed if seed is None else seed)
+
+    def step(
+        self,
+        state: VehicleState,
+        linear_acceleration_world: NDArray,
+        dt: float,
+        *,
+        duration: float = 1.0,
+    ) -> IMUMeasurement:
+        """Generate one measurement from a physical vehicle state.
+
+        The accelerometer convention is specific force
+        ``R.T @ (linear_acceleration + gravity_up)``.
+        """
+        acceleration = np.asarray(linear_acceleration_world, dtype=float)
+        if acceleration.shape != (3,):
+            raise ValueError("linear_acceleration_world must have shape (3,)")
+        rotation = state.rotation_matrix
+        true_accel = rotation.T @ (acceleration + self.cfg.gravity)
+        true_gyro = state.body_rates
+        true_mag = rotation.T @ self.cfg.mag_field_ref
+        temp_factor = 1.0
+        temperature = None
+        if self.temp_model is not None:
+            temperature = self.temp_model.temperature_at(state.time, duration)
+            temp_factor = self.temp_model.noise_scale(temperature)
+            true_accel = true_accel + self.temp_model.accel_offset(temperature)
+            true_gyro = true_gyro + self.temp_model.gyro_offset(temperature)
+            true_mag = true_mag + self.temp_model.mag_offset(temperature)
+        return IMUMeasurement(
+            acceleration=self.accel_model.apply(true_accel, temp_factor, dt),
+            angular_velocity=self.gyro_model.apply(true_gyro, temp_factor, dt),
+            magnetic_field=self.mag_model.apply(true_mag, temp_factor, dt),
+            temperature=temperature,
+        )
+
+    def simulate(self, traj: TrajectoryData) -> IMUData:
+        """Generate sensor readings from a trajectory."""
+        n = len(traj.t)
+        dt = traj.t[1] - traj.t[0] if n > 1 else 0.01
+        duration = traj.t[-1] - traj.t[0]
+
+        accel_out = np.zeros((n, 3))
+        gyro_out = np.zeros((n, 3))
+        mag_out = np.zeros((n, 3))
+        temperature = np.zeros(n) if self.temp_model else None
+
+        for i in range(n):
+            # Rotation matrix from world to body
+            q_xyzw = np.roll(traj.orientation_quat[i], -1)  # wxyz -> xyzw
+            rot = R.from_quat(q_xyzw).as_matrix()
+
+            # --- Accelerometer: gravity in body + linear accel reaction ---
+            gravity_body = rot.T @ self.cfg.gravity
+            lin_accel_body = rot.T @ traj.acceleration[i]
+            true_accel = gravity_body + lin_accel_body
+
+            # --- Gyroscope: angular velocity in body frame ---
+            true_gyro = traj.angular_velocity[i]
+
+            # --- Magnetometer: Earth field in body frame ---
+            true_mag = rot.T @ self.cfg.mag_field_ref
+
+            # Apply temperature effects if enabled
+            temp_factor = 1.0
+            if self.temp_model is not None:
+                temp = self.temp_model.temperature_at(traj.t[i], duration)
+                temperature[i] = temp
+                temp_factor = self.temp_model.noise_scale(temp)
+                true_accel = true_accel + self.temp_model.accel_offset(temp)
+                true_gyro = true_gyro + self.temp_model.gyro_offset(temp)
+                true_mag = true_mag + self.temp_model.mag_offset(temp)
+
+            accel_out[i] = self.accel_model.apply(true_accel, temp_factor, dt)
+            gyro_out[i] = self.gyro_model.apply(true_gyro, temp_factor, dt)
+            mag_out[i] = self.mag_model.apply(true_mag, temp_factor, dt)
+
+        return IMUData(
+            t=traj.t,
+            accel=accel_out,
+            gyro=gyro_out,
+            mag=mag_out,
+            temperature=temperature,
+            accel_bias=self.accel_model.bias.copy(),
+            gyro_bias=self.gyro_model.bias.copy(),
+            mag_bias=self.mag_model.bias.copy(),
+        )

drones_sim/sensors/models.py

@@ -0,0 +1,113 @@
+"""Sensor noise and temperature models for IMU simulation."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+@dataclass
+class SensorNoiseModel:
+    """White noise + bias model for a 3-axis sensor.
+
+    Bias model
+    ----------
+    When ``bias_time_constant`` is finite and ``bias_random_walk_std > 0``,
+    the bias evolves each call to ``apply(dt=...)`` as a discrete
+    first-order Gauss-Markov process:
+
+        b_{k+1} = exp(-dt/tau) * b_k + sigma_b * sqrt(1 - exp(-2dt/tau)) * eta
+
+    where ``tau = bias_time_constant`` and ``sigma_b = bias_random_walk_std``.
+
+    The default ``bias_time_constant=inf`` reproduces the original constant-bias
+    behaviour (backward compatible).  Set ``bias_random_walk_std=0`` to disable
+    the stochastic drive while keeping the exponential decay.
+    """
+    noise_std: float
+    bias_range: float
+    scale_factor_range: tuple[float, float] = (1.0, 1.0)
+    bias_time_constant: float = float('inf')   # tau_b  [s]; inf = constant bias
+    bias_random_walk_std: float = 0.0          # sigma_b [units/s^0.5]
+    rng: np.random.Generator | None = field(default=None, repr=False)
+
+    # Populated on init
+    bias: NDArray = field(init=False)
+    scale_factor: NDArray = field(init=False)
+
+    def __post_init__(self):
+        random = self.rng if self.rng is not None else np.random
+        self.bias = random.uniform(-self.bias_range, self.bias_range, 3)
+        lo, hi = self.scale_factor_range
+        self.scale_factor = random.uniform(lo, hi, 3)
+
+    def apply(
+        self,
+        true_value: NDArray,
+        temp_factor: float = 1.0,
+        dt: float | None = None,
+    ) -> NDArray:
+        """Apply scale factor, bias, and additive noise.
+
+        When ``dt`` is provided and bias dynamics are configured the bias
+        state is updated via the Gauss-Markov recursion before measurement
+        corruption.
+        """
+        # Bias update (Gauss-Markov random walk)
+        if (
+            dt is not None
+            and not np.isinf(self.bias_time_constant)
+            and self.bias_time_constant > 0.0
+        ):
+            decay = np.exp(-dt / self.bias_time_constant)
+            if self.bias_random_walk_std > 0.0:
+                noise_amp = self.bias_random_walk_std * np.sqrt(1.0 - decay ** 2)
+                random = self.rng if self.rng is not None else np.random
+                drive = random.normal(0.0, noise_amp, 3)
+            else:
+                drive = np.zeros(3)
+            self.bias = decay * self.bias + drive
+
+        noisy = true_value * self.scale_factor + self.bias
+        random = self.rng if self.rng is not None else np.random
+        noisy += random.normal(0, self.noise_std * temp_factor, 3)
+        return noisy
+
+
+@dataclass
+class TemperatureModel:
+    """Sinusoidal temperature profile with per-axis sensitivity coefficients."""
+    base_temp: float = 25.0
+    amplitude: float = 10.0
+    rng: np.random.Generator | None = field(default=None, repr=False)
+
+    # Per-axis temperature coefficients (populated randomly)
+    accel_coef: NDArray = field(init=False)
+    gyro_coef: NDArray = field(init=False)
+    mag_coef: NDArray = field(init=False)
+
+    def __post_init__(self):
+        random = self.rng if self.rng is not None else np.random
+        self.accel_coef = random.uniform(-0.002, 0.002, 3)
+        self.gyro_coef = random.uniform(-0.0005, 0.0005, 3)
+        self.mag_coef = random.uniform(-0.05, 0.05, 3)
+
+    def temperature_at(self, t: float, duration: float) -> float:
+        if duration <= 0.0:
+            return self.base_temp
+        return self.base_temp + self.amplitude * np.sin(2 * np.pi * t / duration)
+
+    def noise_scale(self, temp: float) -> float:
+        """Noise increases slightly with temperature deviation."""
+        return 1.0 + 0.01 * (temp - self.base_temp) / 10.0
+
+    def accel_offset(self, temp: float) -> NDArray:
+        return self.accel_coef * (temp - self.base_temp)
+
+    def gyro_offset(self, temp: float) -> NDArray:
+        return self.gyro_coef * (temp - self.base_temp)
+
+    def mag_offset(self, temp: float) -> NDArray:
+        return self.mag_coef * (temp - self.base_temp)

drones_sim/simulation.py

@@ -0,0 +1,283 @@
+"""High-level closed-loop orchestration for reproducible experiments."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+from numpy.typing import NDArray
+
+from .control import GeometricController
+from .dynamics import QuadcopterDynamics
+from .estimation import ExtendedKalmanFilter
+from .math_utils import quat_to_euler
+from .sensors import GPSConfig, GPSSimulator, IMUConfig, IMUSimulator
+from .state import TrajectorySetpoint, VehicleState
+from .trajectory import TrajectoryData
+
+
+@dataclass
+class SimulationConfig:
+    """Timing, sensor, and reproducibility settings."""
+
+    dt: float = 0.01
+    use_estimator: bool = True
+    gps_rate: float = 10.0
+    seed: int = 0
+    # Gravity-as-measurement updates are valid only when translational
+    # acceleration is negligible. IMU acceleration is still used for INS
+    # propagation; opt in to this correction for quasi-static scenarios.
+    correct_accelerometer: bool = False
+    correct_magnetometer: bool = True
+
+    def __post_init__(self) -> None:
+        if self.dt <= 0.0:
+            raise ValueError("dt must be positive")
+        if self.gps_rate <= 0.0:
+            raise ValueError("gps_rate must be positive")
+
+
+@dataclass
+class SimulationResult:
+    """Aligned truth, estimate, reference, actuation, and filter diagnostics."""
+
+    time: NDArray
+    position: NDArray
+    velocity: NDArray
+    quaternion: NDArray
+    body_rates: NDArray
+    estimated_position: NDArray
+    estimated_velocity: NDArray
+    estimated_quaternion: NDArray
+    reference_position: NDArray
+    reference_velocity: NDArray
+    motor_speeds: NDArray
+    commanded_wrench: NDArray
+    allocation_saturated: NDArray
+    metadata: dict[str, object] = field(default_factory=dict)
+
+    @property
+    def tracking_error(self) -> NDArray:
+        return self.position - self.reference_position
+
+    @property
+    def estimation_error(self) -> NDArray:
+        return self.estimated_position - self.position
+
+    def summary(self) -> dict[str, float]:
+        tracking_norm = np.linalg.norm(self.tracking_error, axis=1)
+        estimation_norm = np.linalg.norm(self.estimation_error, axis=1)
+        return {
+            "duration_s": float(self.time[-1] - self.time[0]) if len(self.time) else 0.0,
+            "tracking_rmse_m": float(np.sqrt(np.mean(tracking_norm**2))),
+            "tracking_max_m": float(np.max(tracking_norm)),
+            "estimation_rmse_m": float(np.sqrt(np.mean(estimation_norm**2))),
+            "motor_peak_rad_s": float(np.max(self.motor_speeds)),
+            "allocation_saturation_fraction": float(
+                np.mean(self.allocation_saturated)
+            ),
+        }
+
+
+class ClosedLoopSimulator:
+    """Own the full reference → control → plant → sensors → estimate loop."""
+
+    def __init__(
+        self,
+        vehicle: QuadcopterDynamics | None = None,
+        controller=None,
+        estimator: ExtendedKalmanFilter | None = None,
+        imu: IMUSimulator | None = None,
+        gps: GPSSimulator | None = None,
+        config: SimulationConfig | None = None,
+    ) -> None:
+        self.config = config or SimulationConfig()
+        self.vehicle = vehicle or QuadcopterDynamics(motor_time_constant=0.04)
+        self.controller = controller or GeometricController(self.vehicle)
+        self.estimator = estimator
+        self.imu = imu or IMUSimulator(
+            IMUConfig(
+                accel_scale=(1.0, 1.0),
+                gyro_scale=(1.0, 1.0),
+                mag_scale=(1.0, 1.0),
+            ),
+            seed=self.config.seed,
+        )
+        self.gps = gps or GPSSimulator(
+            GPSConfig(position_noise_std=0.5, update_rate=self.config.gps_rate),
+            seed=self.config.seed + 1,
+        )
+
+    def _reset_estimator(self, initial: VehicleState) -> None:
+        if not self.config.use_estimator:
+            return
+        if self.estimator is None:
+            initial_filter_state = np.zeros(16)
+            initial_filter_state[:3] = initial.position
+            initial_filter_state[3:6] = initial.velocity
+            initial_filter_state[6:10] = initial.quaternion
+            self.estimator = ExtendedKalmanFilter(
+                self.config.dt,
+                initial_state=initial_filter_state,
+                gravity=np.array([0.0, 0.0, self.vehicle.g]),
+                mag_ref=self.imu.cfg.mag_field_ref,
+            )
+        else:
+            self.estimator.x.fill(0.0)
+            self.estimator.x[:3] = initial.position
+            self.estimator.x[3:6] = initial.velocity
+            self.estimator.x[6:10] = initial.quaternion
+            self.estimator.P = np.eye(self.estimator.n) * 0.01
+            self.estimator.P[0:3, 0:3] *= 0.01
+            self.estimator.P[3:6, 3:6] *= 0.1
+            self.estimator.P[6:10, 6:10] *= 0.001
+            self.estimator.P[13:16, 13:16] = np.eye(3) * 0.0001
+            self.estimator.diagnostics = type(self.estimator.diagnostics)()
+            self.estimator._rejection_streak.clear()
+
+    def run(
+        self,
+        reference: TrajectoryData,
+        *,
+        initial_state: VehicleState | None = None,
+    ) -> SimulationResult:
+        """Run one deterministic simulation aligned to ``reference.t``."""
+        if len(reference.t) < 2:
+            raise ValueError("reference must contain at least two samples")
+        sample_dt = np.diff(reference.t)
+        if not np.allclose(sample_dt, self.config.dt, rtol=2e-2, atol=1e-8):
+            raise ValueError(
+                "reference sample period must match SimulationConfig.dt; "
+                f"got median {np.median(sample_dt):.6f}s vs {self.config.dt:.6f}s"
+            )
+        if initial_state is None:
+            initial_state = VehicleState(
+                position=reference.position[0],
+                velocity=reference.velocity[0],
+                quaternion=reference.orientation_quat[0],
+            )
+        self.vehicle.reset(state=initial_state)
+        self.imu.reset(self.config.seed)
+        self.gps.reset(self.config.seed + 1)
+        hover_speed = np.sqrt(
+            self.vehicle.mass * self.vehicle.g / (4.0 * self.vehicle.k_f)
+        )
+        self.vehicle.motor_states[:] = hover_speed
+        self.controller.reset()
+        self._reset_estimator(initial_state)
+
+        count = len(reference.t)
+        shape3 = (count, 3)
+        position = np.zeros(shape3)
+        velocity = np.zeros(shape3)
+        quaternion = np.zeros((count, 4))
+        body_rates = np.zeros(shape3)
+        estimated_position = np.zeros(shape3)
+        estimated_velocity = np.zeros(shape3)
+        estimated_quaternion = np.zeros((count, 4))
+        motor_speeds = np.zeros((count, 4))
+        commanded_wrench = np.zeros((count, 4))
+        allocation_saturated = np.zeros(count, dtype=bool)
+        gps_period = max(1, int(round(1.0 / (self.config.gps_rate * self.config.dt))))
+
+        for index in range(count):
+            true_state = self.vehicle.get_state()
+            measurement = self.imu.step(
+                true_state,
+                self.vehicle.last_acceleration,
+                self.config.dt,
+                duration=float(reference.t[-1]),
+            )
+
+            control_state = true_state
+            if self.config.use_estimator:
+                assert self.estimator is not None
+                self.estimator.predict(
+                    measurement.angular_velocity, measurement.acceleration
+                )
+                if self.config.correct_accelerometer:
+                    self.estimator.correct_accel(measurement.acceleration)
+                if self.config.correct_magnetometer:
+                    self.estimator.correct_mag(measurement.magnetic_field)
+                if index % gps_period == 0:
+                    gps_position, gps_velocity, valid = self.gps.step(
+                        true_state.position, true_state.velocity
+                    )
+                    if valid:
+                        gps_covariance = np.diag(
+                            [self.gps.config.position_noise_std**2] * 3
+                            + [self.gps.config.velocity_noise_std**2] * 3
+                        )
+                        self.estimator.correct_gps(
+                            gps_position, gps_velocity, gps_covariance
+                        )
+                corrected_rates = (
+                    measurement.angular_velocity - self.estimator.x[13:16]
+                )
+                control_state = self.estimator.get_vehicle_state(
+                    corrected_rates, time=true_state.time
+                )
+
+            yaw = float(quat_to_euler(reference.orientation_quat[index])[2])
+            setpoint = TrajectorySetpoint(
+                position=reference.position[index],
+                velocity=reference.velocity[index],
+                acceleration=reference.acceleration[index],
+                yaw=yaw,
+                yaw_rate=float(reference.angular_velocity[index, 2]),
+            )
+            if hasattr(self.controller, "compute_output"):
+                output = self.controller.compute_output(
+                    setpoint, self.config.dt, state=control_state
+                )
+                command = output.motor_speeds
+                commanded_wrench[index] = np.concatenate(
+                    [[output.thrust], output.torque]
+                )
+                allocation_saturated[index] = output.saturated
+            else:
+                previous = reference.position[max(0, index - 1)]
+                command = self.controller.compute(
+                    setpoint.position, setpoint.yaw, self.config.dt, previous
+                )
+                commanded_wrench[index] = (
+                    self.vehicle.allocation_matrix() @ np.square(command)
+                )
+
+            position[index] = true_state.position
+            velocity[index] = true_state.velocity
+            quaternion[index] = true_state.quaternion
+            body_rates[index] = true_state.body_rates
+            estimated_position[index] = control_state.position
+            estimated_velocity[index] = control_state.velocity
+            estimated_quaternion[index] = control_state.quaternion
+            motor_speeds[index] = command
+            self.vehicle.update(self.config.dt, command)
+
+        metadata: dict[str, object] = {
+            "dt": self.config.dt,
+            "controller": type(self.controller).__name__,
+            "estimator": (
+                type(self.estimator).__name__ if self.config.use_estimator else "truth"
+            ),
+            "seed": self.config.seed,
+        }
+        if self.estimator is not None:
+            metadata["ekf_accepted"] = dict(self.estimator.diagnostics.accepted)
+            metadata["ekf_rejected"] = dict(self.estimator.diagnostics.rejected)
+        return SimulationResult(
+            time=reference.t.copy(),
+            position=position,
+            velocity=velocity,
+            quaternion=quaternion,
+            body_rates=body_rates,
+            estimated_position=estimated_position,
+            estimated_velocity=estimated_velocity,
+            estimated_quaternion=estimated_quaternion,
+            reference_position=reference.position.copy(),
+            reference_velocity=reference.velocity.copy(),
+            motor_speeds=motor_speeds,
+            commanded_wrench=commanded_wrench,
+            allocation_saturated=allocation_saturated,
+            metadata=metadata,
+        )