uav-catch/tests/test_dynamics.py

"""Tests for quadcopter dynamics and configuration."""

from __future__ import annotations

from pathlib import Path

import mujoco
import numpy as np
import pytest

from uavcatch.config import EnvironmentConfig, UavConfig, load_control_config, load_environment_config, load_uav_config
from uavcatch.dynamics.model import QuadDynamics
from uavcatch.dynamics.state import QuadState, euler_to_quaternion, normalize_quaternion
from uavcatch.dynamics.wind import WindModel
from uavcatch.sim.mujoco_model import _motor_xy, build_scene_mjcf

ROOT = Path(__file__).resolve().parents[1]


@pytest.fixture
def env_cfg() -> EnvironmentConfig:
  return load_environment_config(ROOT / "environment.toml")


@pytest.fixture
def uav_cfg() -> UavConfig:
  return load_uav_config(ROOT / "uav.toml")


@pytest.fixture
def dynamics(env_cfg: EnvironmentConfig, uav_cfg: UavConfig) -> QuadDynamics:
  return QuadDynamics(env_cfg, uav_cfg)


def test_config_load_defaults(env_cfg: EnvironmentConfig, uav_cfg: UavConfig) -> None:
  assert env_cfg.simulation.dt > 0.0
  assert env_cfg.simulation.integrator == "rk4"
  assert uav_cfg.body.mass == 1.3
  assert uav_cfg.rotor.layout == "x"


def test_control_config_load() -> None:
  cfg = load_control_config(ROOT / "control.toml")
  assert cfg.controller.mode in ("linear", "geometric")
  assert cfg.position_linear.kp == 0.8
  assert cfg.position_geometric.kp == 2.5
  assert cfg.eval.circle.radius == 2.5


def test_config_missing_file() -> None:
  with pytest.raises(FileNotFoundError):
    load_environment_config(ROOT / "missing.toml")


def test_quaternion_normalization_after_integration(dynamics: QuadDynamics) -> None:
  state = QuadState.from_initial(
    position=[0.0, 0.0, 1.0],
    velocity=[0.0, 0.0, 0.0],
    attitude_euler=[0.1, -0.05, 0.2],
    angular_rate=[0.5, -0.3, 0.1],
  )
  omega = np.full(4, dynamics.hover_omega())
  dt = dynamics.env_cfg.simulation.dt

  for _ in range(1000):
    state = dynamics.step(state, omega, dt)
    q_norm = np.linalg.norm(state.quaternion)
    assert abs(q_norm - 1.0) < 1e-6


def test_hover_acceleration_near_zero(dynamics: QuadDynamics) -> None:
  state = QuadState.from_initial(
    position=[0.0, 0.0, 1.0],
    velocity=[0.0, 0.0, 0.0],
    attitude_euler=[0.0, 0.0, 0.0],
    angular_rate=[0.0, 0.0, 0.0],
  )
  omega_hover = dynamics.hover_omega()
  omega = np.full(4, omega_hover)
  deriv = dynamics.derivatives(state, omega)
  assert abs(deriv[5]) < 0.5


def test_free_fall(dynamics: QuadDynamics) -> None:
  state = QuadState.from_initial(
    position=[0.0, 0.0, 5.0],
    velocity=[0.0, 0.0, 0.0],
    attitude_euler=[0.0, 0.0, 0.0],
    angular_rate=[0.0, 0.0, 0.0],
  )
  omega = np.zeros(4)
  dt = dynamics.env_cfg.simulation.dt
  z_prev = state.position[2]

  for _ in range(200):
    state = dynamics.step(state, omega, dt)
    assert state.position[2] <= z_prev + 1e-9
    z_prev = state.position[2]

  assert state.position[2] < 5.0


def test_state_vector_roundtrip() -> None:
  q = euler_to_quaternion(0.1, 0.2, 0.3)
  state = QuadState(
    position=np.array([1.0, 2.0, 3.0]),
    velocity=np.array([0.1, 0.2, 0.3]),
    quaternion=normalize_quaternion(q),
    angular_rate=np.array([0.01, 0.02, 0.03]),
  )
  restored = QuadState.from_vector(state.as_vector())
  np.testing.assert_allclose(restored.position, state.position)
  np.testing.assert_allclose(restored.velocity, state.velocity)
  np.testing.assert_allclose(restored.angular_rate, state.angular_rate)


def test_drag_opposes_motion(dynamics: QuadDynamics) -> None:
  forward = np.array([5.0, 0.0, 0.0], dtype=np.float64)
  force = dynamics.aerodynamic_force(forward, np.zeros(3))
  # Drag must oppose velocity and be purely along -X here.
  assert force[0] < 0.0
  np.testing.assert_allclose(force[1:], 0.0, atol=1e-12)


def test_quadratic_drag_scaling(dynamics: QuadDynamics) -> None:
  f1 = dynamics.aerodynamic_force(np.array([1.0, 0.0, 0.0]), np.zeros(3))
  f2 = dynamics.aerodynamic_force(np.array([2.0, 0.0, 0.0]), np.zeros(3))
  # Quadratic drag: doubling speed quadruples the force (when linear term is 0).
  np.testing.assert_allclose(abs(f2[0]) / abs(f1[0]), 4.0, rtol=1e-6)


def test_drag_uses_relative_wind(dynamics: QuadDynamics) -> None:
  wind = np.array([3.0, 0.0, 0.0], dtype=np.float64)
  # Moving with the wind -> zero relative airspeed -> zero drag.
  force = dynamics.aerodynamic_force(wind.copy(), wind)
  np.testing.assert_allclose(force, 0.0, atol=1e-12)


def test_terminal_velocity_with_drag(dynamics: QuadDynamics) -> None:
  state = QuadState.from_initial(
    position=[0.0, 0.0, 100.0],
    velocity=[0.0, 0.0, 0.0],
    attitude_euler=[0.0, 0.0, 0.0],
    angular_rate=[0.0, 0.0, 0.0],
  )
  omega = np.zeros(4)
  dt = dynamics.env_cfg.simulation.dt
  for _ in range(5000):
    state = dynamics.step(state, omega, dt)
  # With quadratic drag, fall speed must saturate well below the drag-free
  # value (g * t = 9.81 * 10 ~ 98 m/s).
  assert abs(state.velocity[2]) < 60.0


def test_wind_model_gust_oscillates() -> None:
  wind = WindModel(
    base_velocity=np.array([2.0, 0.0, 0.0]),
    gust_amplitude=1.0,
    gust_frequency=1.0,
  )
  assert wind.is_active
  # Quarter period -> peak gust along +X direction.
  np.testing.assert_allclose(wind.sample(0.25)[0], 3.0, atol=1e-6)
  np.testing.assert_allclose(wind.sample(0.0)[0], 2.0, atol=1e-6)


def test_wind_config_parsed(env_cfg: EnvironmentConfig) -> None:
  assert env_cfg.wind.velocity == [0.0, 0.0, 0.0]
  assert hasattr(env_cfg.wind, "gust_amplitude")


def test_mujoco_model_motor_alignment(env_cfg: EnvironmentConfig, uav_cfg: UavConfig) -> None:
  mjcf = build_scene_mjcf(env_cfg, uav_cfg)
  model = mujoco.MjModel.from_xml_string(mjcf)
  arm = uav_cfg.body.arm_length
  expected = _motor_xy(uav_cfg.rotor.layout, arm)

  for idx, (mx, my) in enumerate(expected, start=1):
    body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, f"rotor{idx}")
    assert body_id >= 0
    pos = model.body_pos[body_id]
    np.testing.assert_allclose(pos[0], mx, rtol=0, atol=1e-6)
    np.testing.assert_allclose(pos[1], my, rtol=0, atol=1e-6)