Add doctests

Author: amacati

.github/workflows/testing.yml

@@ -25,3 +25,19 @@ jobs:
         run: |
           pixi --version
           pixi run -e tests tests -v
+
+  test-docs:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Setup Pixi
+        uses: prefix-dev/setup-pixi@v0.9.3
+        with:
+          pixi-version: v0.61.0
+          cache: true
+          cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
+          environments: tests
+          activate-environment: false
+          locked: true
+      - name: Run doctests
+        run: pixi run -e tests test-docs

docs/examples/index.md

@@ -17,10 +17,11 @@ pos     = np.zeros(3)
 quat    = np.array([0., 0., 0., 1.])   # xyzw — identity (no rotation)
 vel     = np.zeros(3)
 ang_vel = np.zeros(3)
+rotor_vel = np.ones(4) * 12_000.
 cmd     = np.full(4, 15_000.)           # motor RPMs
 
 pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = model(
-    pos, quat, vel, ang_vel, cmd
+    pos, quat, vel, ang_vel, cmd, rotor_vel
 )
 ```
 
@@ -31,23 +32,40 @@ The outputs are the time derivatives of each state variable — the right-hand s
 In the previous example `rotor_vel_dot` is `None` because we didn't tell the model what speed the motors are currently at. The model just assumed they were already at the commanded RPM. That's a reasonable approximation for slow maneuvers, but real motors take time to spin up and down. Passing `rotor_vel` enables the rotor dynamics, which computes the acceleration of each motor toward its target.
 
 ```python
+import numpy as np
+from drone_models import parametrize
+from drone_models.first_principles import dynamics
+
+model = parametrize(dynamics, drone_model="cf2x_L250")
+pos     = np.zeros(3)
+quat    = np.array([0., 0., 0., 1.])
+vel     = np.zeros(3)
+ang_vel = np.zeros(3)
+cmd     = np.full(4, 15_000.)
 # The motors are at 12 000 RPM but commanded to 15 000 — they're spinning up.
 rotor_vel = np.full(4, 12_000.)
 
 pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = model(
     pos, quat, vel, ang_vel, cmd, rotor_vel=rotor_vel
 )
-print(rotor_vel_dot)  # positive — rotors accelerating toward cmd
+rotor_vel_dot  # positive — rotors accelerating toward cmd
 ```
 
 ## Fitted models
 
 The first-principles model requires individual motor RPMs as input, which means you need rotor-level commands. The fitted models — `so_rpy`, `so_rpy_rotor`, `so_rpy_rotor_drag` — take a higher-level command instead: roll, pitch, yaw setpoints in radians plus collective thrust in Newtons. This matches the command interface of typical flight controllers and makes them convenient for control design and system identification.
 
 ```python
-from drone_models.so_rpy_rotor_drag import dynamics as srrd
+import numpy as np
+from drone_models import parametrize
+from drone_models.so_rpy import dynamics
 
-model = parametrize(srrd, drone_model="cf2x_L250")
+pos     = np.zeros(3)
+quat    = np.array([0., 0., 0., 1.])
+vel     = np.zeros(3)
+ang_vel = np.zeros(3)
+
+model = parametrize(dynamics, drone_model="cf2x_L250")
 
 # Collective thrust near hover: mass * g ≈ 0.0319 * 9.81 ≈ 0.31 N
 cmd = np.array([0., 0., 0., 0.31])   # [roll_rad, pitch_rad, yaw_rad, thrust_N]
@@ -73,10 +91,11 @@ pos     = torch.zeros(3)
 quat    = torch.tensor([0., 0., 0., 1.])
 vel     = torch.zeros(3)
 ang_vel = torch.zeros(3)
+rotor_vel = torch.ones(4) * 12_000.
 cmd     = torch.full((4,), 15_000.)
 
 pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = model(
-    pos, quat, vel, ang_vel, cmd
+    pos, quat, vel, ang_vel, cmd, rotor_vel
 )
 ```
 
@@ -85,6 +104,14 @@ pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = model(
 The same model handles arbitrary leading batch dimensions — no special API, no loops. Add a leading dimension to all state and command arrays and the model evaluates all instances in a single call. This works identically across all backends.
 
 ```python
+import torch
+
+from drone_models import parametrize
+from drone_models.first_principles import dynamics
+
+# Parameters are stored as torch tensors — no per-call conversion needed.
+model = parametrize(dynamics, drone_model="cf2x_L250", xp=torch)
+
 N = 1_000
 
 pos       = torch.zeros(N, 3)
@@ -97,7 +124,7 @@ rotor_vel = torch.full((N, 4), 15_000.)
 pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = model(
     pos, quat, vel, ang_vel, cmd, rotor_vel=rotor_vel
 )
-print(vel_dot.shape)   # (1000, 3)
+vel_dot.shape   # (1000, 3)
 ```
 
 ## Overriding parameters at call time
@@ -117,12 +144,13 @@ pos     = jnp.zeros(3)
 quat    = jnp.array([0., 0., 0., 1.])
 vel     = jnp.zeros(3)
 ang_vel = jnp.zeros(3)
+rotor_vel = jnp.ones(4) * 12_000.
 cmd     = jnp.full((4,), 15_000.)
 
 # The model was parametrized with mass=0.0319 kg.
 # Simulate the same drone carrying a 10 g payload for this one call:
 pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = model(
-    pos, quat, vel, ang_vel, cmd, mass=jnp.float32(0.0419)
+    pos, quat, vel, ang_vel, cmd, rotor_vel, mass=jnp.float32(0.0419)
 )
 ```
 
@@ -164,7 +192,7 @@ nominal_mass = model.keywords["mass"]   # scalar
 nominal_J    = model.keywords["J"]      # (3, 3)
 
 key, k1, k2 = jax.random.split(key, 3)
-mass_batch  = nominal_mass * jax.random.uniform(k1, (N,),      minval=0.9, maxval=1.1)
+mass_batch  = nominal_mass * jax.random.uniform(k1, (N, 1),      minval=0.9, maxval=1.1)
 J_batch     = nominal_J    * jax.random.uniform(k2, (N, 3, 3), minval=0.9, maxval=1.1)
 J_inv_batch = jnp.linalg.inv(J_batch)
 
@@ -173,5 +201,5 @@ pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = step(
     pos, quat, vel, ang_vel, cmd, rotor_vel,
     mass_batch, J_batch, J_inv_batch,
 )
-print(vel_dot.shape)   # (4096, 3)
+vel_dot.shape   # (4096, 3)
 ```

docs/get-started/installation.md

@@ -88,6 +88,5 @@ If your drone is not listed, you can fit parameters from your own flight data us
 
 ```python
 from drone_models import available_models
-print(list(available_models))
-# ['first_principles', 'so_rpy', 'so_rpy_rotor', 'so_rpy_rotor_drag']
+list(available_models)  # ['first_principles', 'so_rpy', 'so_rpy_rotor', 'so_rpy_rotor_drag']
 ```

docs/get-started/quick-start.md

@@ -33,13 +33,14 @@ model = parametrize(dynamics, drone_model="cf2x_L250")
 pos     = np.zeros(3)                   # [m]
 quat    = np.array([0., 0., 0., 1.])   # xyzw — identity (no rotation)
 vel     = np.zeros(3)                   # [m/s]
+rotor_vel = np.ones(4) * 12_000.        # [RPM] — motors are spinning but not yet at the 15 000 RPM
 ang_vel = np.zeros(3)                   # [rad/s]
 
 # Command: all four motors at 15 000 RPM (rough hover point for cf2x_L250).
 cmd = np.full(4, 15_000.)              # [RPM]
 
 pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = model(
-    pos, quat, vel, ang_vel, cmd
+    pos, quat, vel, ang_vel, cmd, rotor_vel
 )
 ```
 
@@ -64,13 +65,23 @@ These are the right-hand side of the continuous-time ODE $\dot{x} = f(x, u)$. To
 Real motors don't respond instantaneously to commands. Passing the current motor state as `rotor_vel` enables the rotor dynamics model, which computes how the motors accelerate or decelerate toward the commanded RPM.
 
 ```python
+import numpy as np
+from drone_models import parametrize
+from drone_models.first_principles import dynamics
+
+model = parametrize(dynamics, drone_model="cf2x_L250")
+pos     = np.zeros(3)
+quat    = np.array([0., 0., 0., 1.])
+vel     = np.zeros(3)
+ang_vel = np.zeros(3)
+cmd     = np.full(4, 15_000.)
 # Current RPMs lag behind the 15 000 RPM command — motors are still spinning up.
 rotor_vel = np.full(4, 12_000.)
 
 pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = model(
     pos, quat, vel, ang_vel, cmd, rotor_vel=rotor_vel
 )
-print(rotor_vel_dot)  # positive — rotors accelerating toward cmd
+rotor_vel_dot  # positive — rotors accelerating toward cmd
 ```
 
 ## Next steps

docs/index.md

@@ -17,10 +17,11 @@ pos     = np.zeros(3)
 quat    = np.array([0., 0., 0., 1.])   # xyzw, identity
 vel     = np.zeros(3)
 ang_vel = np.zeros(3)
+rotor_vel = np.ones(4) * 12_000.
 cmd     = np.full(4, 15_000.)           # motor RPMs, near hover
 
 pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = model(
-    pos, quat, vel, ang_vel, cmd
+    pos, quat, vel, ang_vel, cmd, rotor_vel
 )
 ```
 

docs/user-guide/batching.md

@@ -20,7 +20,7 @@ rotor_vel = jnp.full((N, 4), 15_000.)
 pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = model(
     pos, quat, vel, ang_vel, cmd, rotor_vel=rotor_vel
 )
-print(vel_dot.shape)  # (1000, 3)
+vel_dot.shape  # (1000, 3)
 ```
 
 A runnable version of this example is in [Examples: Batched evaluation](../examples/index.md#batched-evaluation).
@@ -30,15 +30,21 @@ A runnable version of this example is in [Examples: Batched evaluation](../examp
 Any number of leading dimensions works. A common pattern is a grid of environments, each containing multiple drones:
 
 ```python
+import jax.numpy as jnp
+from drone_models import parametrize
+from drone_models.first_principles import dynamics
+
+model = parametrize(dynamics, drone_model="cf2x_L250", xp=jnp)
 # 50 environments, 20 drones each
 pos     = jnp.zeros((50, 20, 3))
 quat    = jnp.broadcast_to(jnp.array([0., 0., 0., 1.]), (50, 20, 4))
 vel     = jnp.zeros((50, 20, 3))
 ang_vel = jnp.zeros((50, 20, 3))
+rotor_vel = jnp.full((50, 20, 4), 12_000.)
 cmd     = jnp.full((50, 20, 4), 15_000.)
 
-vel_dot, *_ = model(pos, quat, vel, ang_vel, cmd)
-print(vel_dot.shape)  # (50, 20, 3)
+vel_dot, *_ = model(pos, quat, vel, ang_vel, cmd, rotor_vel)
+vel_dot.shape  # (50, 20, 3)
 ```
 
 ## Domain randomization
@@ -56,6 +62,10 @@ from drone_models.first_principles import dynamics
 N   = 4_096
 key = jax.random.PRNGKey(0)
 
+pos, vel, ang_vel = jnp.zeros((N, 3)), jnp.zeros((N, 3)), jnp.zeros((N, 3))
+quat = jnp.tile(jnp.array([0., 0., 0., 1.]), (N, 1))
+cmd = jnp.full((N, 4), 15_000.)
+rotor_vel = jnp.full((N, 4), 15_000.)
 model = parametrize(dynamics, drone_model="cf2x_L250", xp=jnp)
 nominal_mass = model.keywords["mass"]
 nominal_J    = model.keywords["J"]
@@ -68,7 +78,7 @@ def step(pos, quat, vel, ang_vel, cmd, rotor_vel, mass, J, J_inv):
     )
 
 key, k1, k2 = jax.random.split(key, 3)
-mass_batch  = nominal_mass * jax.random.uniform(k1, (N,),      minval=0.9, maxval=1.1)
+mass_batch  = nominal_mass * jax.random.uniform(k1, (N, 1),      minval=0.9, maxval=1.1)
 J_batch     = nominal_J    * jax.random.uniform(k2, (N, 3, 3), minval=0.9, maxval=1.1)
 J_inv_batch = jnp.linalg.inv(J_batch)
 
@@ -78,7 +88,7 @@ vel_dot = step(pos, quat, vel, ang_vel, cmd, rotor_vel,
 
 **Option 2 — mutate `model.keywords` directly.** Simpler when you don't need JIT or are happy to retrace. Replace a scalar parameter with a `(N,)` array and each element in the batch uses its own value.
 
-```python
+```{ .python notest }
 model.keywords["mass"] = nominal_mass * mass_batch  # shape (N,)
 vel_dot = model(pos, quat, vel, ang_vel, cmd)[2]
 ```

docs/user-guide/models.md

@@ -27,11 +27,17 @@ The full rigid-body physics model. The command is four individual motor RPMs. Th
 Working at the rotor-velocity level means you need a controller that converts higher-level commands — position setpoints, attitude + collective thrust — down to individual motor RPMs. [drone-controllers](https://utiasdsl.github.io/drone-controllers/) provides a matching set of controllers designed for exactly this interface.
 
 ```python
+import numpy as np
 from drone_models import parametrize
 from drone_models.first_principles import dynamics
 
 model = parametrize(dynamics, drone_model="cf2x_L250")
 
+pos, vel, ang_vel = np.zeros((3,)), np.zeros((3,)), np.zeros((3,))
+quat = np.array([0., 0., 0., 1.])
+cmd = np.full((4,), 15_000.)
+rotor_vel = np.full((4,), 12_000.)
+
 pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot = model(
     pos, quat, vel, ang_vel,
     cmd,        # shape (4,) — motor RPMs
@@ -45,7 +51,7 @@ See the [`first_principles` API reference](../reference/drone_models/first_princ
 
 A fitted second-order model where the command is `[roll_rad, pitch_rad, yaw_rad, thrust_N]` — the same interface used by most flight controller firmware. First-order thrust dynamics model motor spin-up delay, and a linear body-frame drag term accounts for aerodynamic resistance. All coefficients are identified from flight data rather than derived from physics, which makes the model easy to calibrate and well-suited to real-time control.
 
-```python
+```{ .python notest }
 from drone_models.so_rpy_rotor_drag import dynamics
 
 model = parametrize(dynamics, drone_model="cf2x_L250")
@@ -78,6 +84,14 @@ from drone_models.so_rpy import dynamics
 All four models accept optional `dist_f` (external force, world frame, N) and `dist_t` (external torque, body frame, N·m) arguments. These are useful for modelling wind, contact forces, or other perturbations without modifying the model itself.
 
 ```python
+import numpy as np
+from drone_models import parametrize
+from drone_models.so_rpy import dynamics
+
+model = parametrize(dynamics, drone_model="cf2x_L250")
+pos, vel, ang_vel = np.zeros((3,)), np.zeros((3,)), np.zeros((3,))
+quat = np.array([0., 0., 0., 1.])
+cmd = np.array([0., 0., 0., 0.31])
 dist_f = np.array([0.05, 0., 0.])  # 50 mN headwind [N]
 dist_t = np.zeros(3)
 
@@ -95,8 +109,8 @@ from drone_models import model_features
 from drone_models.first_principles import dynamics as fp
 from drone_models.so_rpy import dynamics as srpy
 
-print(model_features(fp))     # {'rotor_dynamics': True}
-print(model_features(srpy))   # {'rotor_dynamics': False}
+model_features(fp)     # {'rotor_dynamics': True}
+model_features(srpy)   # {'rotor_dynamics': False}
 ```
 
 ---