From ad7ed8157dedbacb4fb30841bd3b7f08d5a5232b Mon Sep 17 00:00:00 2001
From: radu_dell <radu141@gmail.com>
Date: Tue, 22 Sep 2026 18:28:32 +0200
Subject: [PATCH 01/10] Adds downwash implementation from the ETH paper with
 constant power assumption

---
 examples/plugins/downwash.py | 139 +++++++++++++++++++++++++++++++++++
 1 file changed, 139 insertions(+)
 create mode 100644 examples/plugins/downwash.py

diff --git a/examples/plugins/downwash.py b/examples/plugins/downwash.py
new file mode 100644
index 00000000..1f2e34e5
--- /dev/null
+++ b/examples/plugins/downwash.py
@@ -0,0 +1,139 @@
+"""Minimal far-field downwash thrust-loss plugin.
+
+This models the downwash of level, hovering, identical Crazyflies using the
+far-field jet from Bauersfeld et al. (arXiv:2403.13321). It belongs after
+``force_torque_controller`` and before ``clip_rotor_vel_cmd``.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from crazyflow.control.mellinger import force_torque2rotor_vel
+from crazyflow.sim.pipeline import insert_fn_before
+
+if TYPE_CHECKING:
+    from crazyflow.sim import Sim
+    from crazyflow.sim.data import SimData
+
+# Physical parameters for the cf21B_500
+AIR_DENSITY = 1.225  # kg/m^3
+PROPELLER_RADIUS = 27.5e-3  # m
+MOTOR_DISTANCE = 0.035355*2  # m, distance between opposite motors 
+N_PROPELLERS = 4
+GRAVITY = 9.81
+
+# Far-field fit in Eq. (9) of the paper.
+BD = 10.11
+S = 0.07668
+S0 = -5.817
+
+
+def _thrust_loss(alpha: jax.Array) -> jax.Array:
+    """Return T_effective / T_nominal at constant aerodynamic power.
+
+    Equating the still-air and downwash momentum-theory powers yields the
+    positive root of eta**3 + alpha * eta - 1 = 0, with alpha = U_D / U_H.
+    """
+    lower = jnp.zeros_like(alpha)
+    upper = jnp.ones_like(alpha)
+
+    def bisect(_: int, bounds: tuple[jax.Array, jax.Array]) -> tuple[jax.Array, jax.Array]:
+        lower, upper = bounds
+        eta = (lower + upper) / 2.0
+        above_root = eta**3 + alpha * eta > 1.0
+        return jnp.where(above_root, lower, eta), jnp.where(above_root, eta, upper)
+
+    lower, upper = jax.lax.fori_loop(0, 24, bisect, (lower, upper))
+    return (lower + upper) / 2.0
+
+
+def downwash_fn(data: SimData) -> SimData:
+    """Reduce every lower drone's nominal rotor force by its downwash loss."""
+    # Axis 1 indexes the source drone; axis 2 indexes the target drone.
+    source_to_target = data.states.pos[:, :, None, :] - data.states.pos[:, None, :, :]
+    s = source_to_target[..., 2]  # Positive only for targets below a source.
+    r = jnp.linalg.vector_norm(source_to_target[..., :2], axis=-1)
+    s_normalized = s / MOTOR_DISTANCE
+
+    mass = data.params.mass[0] 
+    u_hover = jnp.sqrt(
+        mass * GRAVITY
+        / (2.0 * AIR_DENSITY * jnp.pi * PROPELLER_RADIUS**2 * N_PROPELLERS)
+    )
+    half_width = S * (s_normalized - S0)
+    centerline_velocity = u_hover * BD / (s_normalized - S0)
+    xi = (r / MOTOR_DISTANCE) / half_width
+    u_downwash = centerline_velocity / (1.0 + (jnp.sqrt(2.0) - 1.0) * xi**2) ** 2
+    u_downwash = jnp.where(s_normalized > 2.5, u_downwash, 0.0)
+    u_downwash = jnp.sum(u_downwash, axis=1)  # Sum all sources at each target.
+
+    eta = _thrust_loss(u_downwash / u_hover)
+    force_torque = data.controls.force_torque
+
+    # Scale only thrust
+    effective_thrust = force_torque.cmd[..., 0] * eta
+
+    rotor_vel = force_torque2rotor_vel(
+        effective_thrust[..., None], force_torque.cmd[..., 1:], **force_torque.params
+    )
+    return data.replace(controls=data.controls.replace(rotor_vel=rotor_vel))
+
+
+def install_downwash(sim: Sim) -> None:
+    """Add downwash after allocation and rebuild the compiled step function."""
+    insert_fn_before(sim.step_pipeline, "clip_rotor_vel_cmd", downwash_fn)
+    sim.build_step_fn()
+
+
+def main(plot: bool = True) -> None:
+    """Hover drone 0 while drone 1 flies straight through its downwash."""
+    from crazyflow.sim import Sim
+
+    sim = Sim(n_drones=2, drone="cf21B_500", control="state")
+    install_downwash(sim)
+
+    upper_pos = np.array([0.0, 0.0, 1.2])
+    lower_start = np.array([-0.5, 0.0, 1.0])
+    sim.data = sim.data.replace(
+        states=sim.data.states.replace(pos=jnp.array([[upper_pos, lower_start]]))
+    )
+    sim.build_default_data()
+
+    duration = 3.0
+    speed = 1.0 / duration
+    command = np.zeros((1, 2, 16))
+    command[..., 9:13] = [0.0, 0.0, 0.0, 1.0]  # level quaternion (xyzw)
+    command[0, 0, :3] = upper_pos
+    z_positions = []
+
+    for step in range(int(duration * sim.control_freq)):
+        t = step / sim.control_freq
+        command[0, 1, :3] = [-0.5 + speed * t, 0.0, 1.0]
+        command[0, 1, 3:6] = [speed, 0.0, 0.0]
+        sim.state_control(command)
+        sim.step(sim.freq // sim.control_freq)
+        z_positions.append(np.asarray(sim.data.states.pos[0, :, 2]))
+        sim.render()
+
+    sim.close()
+    if plot:
+        import matplotlib.pyplot as plt
+
+        t = np.arange(len(z_positions)) / sim.control_freq
+        z_positions = np.asarray(z_positions)
+        plt.plot(t, z_positions[:, 0], label="upper drone")
+        plt.plot(t, z_positions[:, 1], label="lower drone")
+        plt.xlabel("Time (s)")
+        plt.ylabel("z position (m)")
+        plt.legend()
+        plt.show()
+
+
+if __name__ == "__main__":
+    main()
+    

From a9e7a5d5265362a319e0b1adbe6d6c785ef90277 Mon Sep 17 00:00:00 2001
From: radu_workstation <radu141@gmail.com>
Date: Wed, 23 Sep 2026 17:55:21 +0200
Subject: [PATCH 02/10] Reformulates downwash as force on the individual rotor
 blades. Generates torques as well

---
 examples/plugins/downwash.py | 156 +++++++++++++++++++++--------------
 1 file changed, 94 insertions(+), 62 deletions(-)

diff --git a/examples/plugins/downwash.py b/examples/plugins/downwash.py
index 1f2e34e5..4c71ade6 100644
--- a/examples/plugins/downwash.py
+++ b/examples/plugins/downwash.py
@@ -1,123 +1,147 @@
-"""Minimal far-field downwash thrust-loss plugin.
+"""Minimal far-field downwash external-wrench plugin.
 
 This models the downwash of level, hovering, identical Crazyflies using the
-far-field jet from Bauersfeld et al. (arXiv:2403.13321). It belongs after
-``force_torque_controller`` and before ``clip_rotor_vel_cmd``.
+far-field jet from Bauersfeld et al. (arXiv:2403.13321) and the thrust-decay
+model of Su et al. (arXiv:2207.09645).
 """
 
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-import jax
 import jax.numpy as jnp
 import numpy as np
+from jax.scipy.spatial.transform import Rotation as R
 
-from crazyflow.control.mellinger import force_torque2rotor_vel
+from crazyflow.control.transform import motor_force2rotor_vel
+from crazyflow.sim import Sim
 from crazyflow.sim.pipeline import insert_fn_before
 
 if TYPE_CHECKING:
-    from crazyflow.sim import Sim
     from crazyflow.sim.data import SimData
 
 # Physical parameters for the cf21B_500
 AIR_DENSITY = 1.225  # kg/m^3
 PROPELLER_RADIUS = 27.5e-3  # m
-MOTOR_DISTANCE = 0.035355*2  # m, distance between opposite motors 
+MOTOR_DISTANCE = 0.035355 * 2  # m, distance between opposite motors
 N_PROPELLERS = 4
 GRAVITY = 9.81
 
-# Far-field fit in Eq. (9) of the paper.
+# This must be fitted for the propeller/downwash setup.
+THRUST_DECAY_COEFFICIENT = 0.05  # s/m
+
+# Far-field fit in Eq. (9) of the Bauersfeld paper.
 BD = 10.11
 S = 0.07668
 S0 = -5.817
 
 
-def _thrust_loss(alpha: jax.Array) -> jax.Array:
-    """Return T_effective / T_nominal at constant aerodynamic power.
+def downwash_fn(data: SimData) -> SimData:
+    """Apply downwash-induced thrust loss as a world-frame external wrench.
 
-    Equating the still-air and downwash momentum-theory powers yields the
-    positive root of eta**3 + alpha * eta - 1 = 0, with alpha = U_D / U_H.
+    The source flow originates at each drone centre, while the field is sampled
+    at every target rotor.
     """
-    lower = jnp.zeros_like(alpha)
-    upper = jnp.ones_like(alpha)
-
-    def bisect(_: int, bounds: tuple[jax.Array, jax.Array]) -> tuple[jax.Array, jax.Array]:
-        lower, upper = bounds
-        eta = (lower + upper) / 2.0
-        above_root = eta**3 + alpha * eta > 1.0
-        return jnp.where(above_root, lower, eta), jnp.where(above_root, eta, upper)
-
-    lower, upper = jax.lax.fori_loop(0, 24, bisect, (lower, upper))
-    return (lower + upper) / 2.0
-
+    rotation = R.from_quat(data.states.quat)
+    rotor_offsets = data.params.L * jnp.array(
+        [[1.0, -1.0, 0.0], [-1.0, -1.0, 0.0], [-1.0, 1.0, 0.0], [1.0, 1.0, 0.0]]
+    )
+    rotor_offsets_world = jnp.swapaxes(rotation.as_matrix() @ rotor_offsets.T, -1, -2)
+    rotor_positions = data.states.pos[..., None, :] + rotor_offsets_world
 
-def downwash_fn(data: SimData) -> SimData:
-    """Reduce every lower drone's nominal rotor force by its downwash loss."""
-    # Axis 1 indexes the source drone; axis 2 indexes the target drone.
-    source_to_target = data.states.pos[:, :, None, :] - data.states.pos[:, None, :, :]
+    # Axis 1 indexes the source drone, axis 2 the target, and axis 3 its rotor.
+    source_to_target = data.states.pos[:, :, None, None, :] - rotor_positions[:, None, :, :, :]
     s = source_to_target[..., 2]  # Positive only for targets below a source.
     r = jnp.linalg.vector_norm(source_to_target[..., :2], axis=-1)
     s_normalized = s / MOTOR_DISTANCE
 
-    mass = data.params.mass[0] 
+    mass = data.params.mass[0]
     u_hover = jnp.sqrt(
-        mass * GRAVITY
-        / (2.0 * AIR_DENSITY * jnp.pi * PROPELLER_RADIUS**2 * N_PROPELLERS)
+        mass * GRAVITY / (2.0 * AIR_DENSITY * jnp.pi * PROPELLER_RADIUS**2 * N_PROPELLERS)
     )
     half_width = S * (s_normalized - S0)
     centerline_velocity = u_hover * BD / (s_normalized - S0)
     xi = (r / MOTOR_DISTANCE) / half_width
     u_downwash = centerline_velocity / (1.0 + (jnp.sqrt(2.0) - 1.0) * xi**2) ** 2
-    u_downwash = jnp.where(s_normalized > 2.5, u_downwash, 0.0)
-    u_downwash = jnp.sum(u_downwash, axis=1)  # Sum all sources at each target.
-
-    eta = _thrust_loss(u_downwash / u_hover)
-    force_torque = data.controls.force_torque
-
-    # Scale only thrust
-    effective_thrust = force_torque.cmd[..., 0] * eta
-
-    rotor_vel = force_torque2rotor_vel(
-        effective_thrust[..., None], force_torque.cmd[..., 1:], **force_torque.params
+    # u_downwash = jnp.where(s_normalized > 2.5, u_downwash, 0.0)
+    u_downwash = jnp.sum(u_downwash, axis=1)  # Sum all sources at each target rotor.
+
+    # Eq. (5) in Su et al.: each motor loses a fraction b_v * U_D of its
+    # current thrust.  Clamp this extrapolation so effective thrust is never
+    # negative outside the fitted range.
+    loss_fraction = jnp.clip(THRUST_DECAY_COEFFICIENT * u_downwash, 0.0, 1.0)
+    rotor_vel = data.states.rotor_vel
+    k0, k1, k2 = (
+        data.params.rpm2thrust[..., 0],
+        data.params.rpm2thrust[..., 1],
+        data.params.rpm2thrust[..., 2],
     )
-    return data.replace(controls=data.controls.replace(rotor_vel=rotor_vel))
-
-
-def install_downwash(sim: Sim) -> None:
-    """Add downwash after allocation and rebuild the compiled step function."""
-    insert_fn_before(sim.step_pipeline, "clip_rotor_vel_cmd", downwash_fn)
-    sim.build_step_fn()
+    motor_thrust = k0 + k1 * rotor_vel + k2 * rotor_vel**2
+    thrust_delta = -loss_fraction * motor_thrust
+
+    # Map the per-motor force changes to a body-frame wrench, as in Eq. (7).
+    total_thrust_delta = jnp.sum(thrust_delta, axis=-1)
+    zeros = jnp.zeros_like(total_thrust_delta)
+    force_body = jnp.stack((zeros, zeros, total_thrust_delta), axis=-1)
+
+    lever = jnp.array([1.0, 1.0, 0.0])
+    torque_body = (data.params.mixing_matrix @ (thrust_delta * data.params.L)[..., None])[
+        ..., 0
+    ] * lever
+
+    # Account for the corresponding reaction-torque change about body z.
+    effective_motor_thrust = jnp.maximum(motor_thrust + thrust_delta, 0.0)
+    effective_rotor_vel = motor_force2rotor_vel(effective_motor_thrust, data.params.rpm2thrust)
+    c0, c1, c2 = (
+        data.params.rpm2torque[..., 0],
+        data.params.rpm2torque[..., 1],
+        data.params.rpm2torque[..., 2],
+    )
+    motor_torque = c0 + c1 * rotor_vel + c2 * rotor_vel**2
+    effective_motor_torque = c0 + c1 * effective_rotor_vel + c2 * effective_rotor_vel**2
+    reaction_torque_delta = effective_motor_torque - motor_torque
+    torque_body = torque_body + (data.params.mixing_matrix @ reaction_torque_delta[..., None])[
+        ..., 0
+    ] * jnp.array([0.0, 0.0, 1.0])
+
+    states = data.states.replace(
+        force=rotation.apply(force_body), torque=rotation.apply(torque_body)
+    )
+    return data.replace(states=states)
 
 
 def main(plot: bool = True) -> None:
     """Hover drone 0 while drone 1 flies straight through its downwash."""
-    from crazyflow.sim import Sim
-
     sim = Sim(n_drones=2, drone="cf21B_500", control="state")
-    install_downwash(sim)
 
-    upper_pos = np.array([0.0, 0.0, 1.2])
-    lower_start = np.array([-0.5, 0.0, 1.0])
+    insert_fn_before(sim.step_pipeline, "integration", downwash_fn)
+    sim.build_step_fn()
+
+    upper_pos = np.array([0.0, 0.0, 1.0])
+    lower_start = np.array([-0.5, 0.0, 0.5])
     sim.data = sim.data.replace(
         states=sim.data.states.replace(pos=jnp.array([[upper_pos, lower_start]]))
     )
     sim.build_default_data()
 
-    duration = 3.0
+    duration = 6.0
     speed = 1.0 / duration
     command = np.zeros((1, 2, 16))
     command[..., 9:13] = [0.0, 0.0, 0.0, 1.0]  # level quaternion (xyzw)
     command[0, 0, :3] = upper_pos
     z_positions = []
+    downwash_force_z = []
+    downwash_pitch_torque = []
 
     for step in range(int(duration * sim.control_freq)):
         t = step / sim.control_freq
-        command[0, 1, :3] = [-0.5 + speed * t, 0.0, 1.0]
+        command[0, 1, :3] = [-0.5 + speed * t, 0.0, 0.5]
         command[0, 1, 3:6] = [speed, 0.0, 0.0]
         sim.state_control(command)
         sim.step(sim.freq // sim.control_freq)
         z_positions.append(np.asarray(sim.data.states.pos[0, :, 2]))
+        downwash_force_z.append(np.asarray(sim.data.states.force[0, 1, 2]))
+        downwash_pitch_torque.append(np.asarray(sim.data.states.torque[0, 1, 1]))
         sim.render()
 
     sim.close()
@@ -126,14 +150,22 @@ def main(plot: bool = True) -> None:
 
         t = np.arange(len(z_positions)) / sim.control_freq
         z_positions = np.asarray(z_positions)
-        plt.plot(t, z_positions[:, 0], label="upper drone")
-        plt.plot(t, z_positions[:, 1], label="lower drone")
-        plt.xlabel("Time (s)")
-        plt.ylabel("z position (m)")
-        plt.legend()
+        fig, axes = plt.subplots(3, 1, sharex=True)
+        axes[0].plot(t, z_positions[:, 0], label="upper drone")
+        axes[0].plot(t, z_positions[:, 1], label="lower drone")
+        axes[0].set_ylabel("z position (m)")
+        axes[0].legend()
+
+        axes[1].plot(t, downwash_force_z, label="lower drone")
+        axes[1].set_ylabel("downwash force z (N)")
+        axes[1].legend()
+
+        axes[2].plot(t, downwash_pitch_torque, label="lower drone")
+        axes[2].set_xlabel("Time (s)")
+        axes[2].set_ylabel("downwash pitch torque y (Nm)")
+        axes[2].legend()
         plt.show()
 
 
 if __name__ == "__main__":
     main()
-    

From baef4efb48c9b403e754aa87a09e491de3d0e378 Mon Sep 17 00:00:00 2001
From: radu_workstation <radu141@gmail.com>
Date: Thu, 24 Sep 2026 11:16:27 +0200
Subject: [PATCH 03/10] Adds plotting for downwash heatmap

---
 examples/plugins/downwash.py | 114 +++++++++++++++++++++++++++++------
 1 file changed, 96 insertions(+), 18 deletions(-)

diff --git a/examples/plugins/downwash.py b/examples/plugins/downwash.py
index 4c71ade6..991fb236 100644
--- a/examples/plugins/downwash.py
+++ b/examples/plugins/downwash.py
@@ -28,7 +28,7 @@
 GRAVITY = 9.81
 
 # This must be fitted for the propeller/downwash setup.
-THRUST_DECAY_COEFFICIENT = 0.05  # s/m
+THRUST_DECAY_COEFFICIENT = 0.07 # s/m
 
 # Far-field fit in Eq. (9) of the Bauersfeld paper.
 BD = 10.11
@@ -63,13 +63,13 @@ def downwash_fn(data: SimData) -> SimData:
     centerline_velocity = u_hover * BD / (s_normalized - S0)
     xi = (r / MOTOR_DISTANCE) / half_width
     u_downwash = centerline_velocity / (1.0 + (jnp.sqrt(2.0) - 1.0) * xi**2) ** 2
-    # u_downwash = jnp.where(s_normalized > 2.5, u_downwash, 0.0)
+
+    # This prevents "negative" downwash
+    u_downwash = jnp.where(s_normalized > 0.1, u_downwash, 0.0)
     u_downwash = jnp.sum(u_downwash, axis=1)  # Sum all sources at each target rotor.
 
-    # Eq. (5) in Su et al.: each motor loses a fraction b_v * U_D of its
-    # current thrust.  Clamp this extrapolation so effective thrust is never
-    # negative outside the fitted range.
-    loss_fraction = jnp.clip(THRUST_DECAY_COEFFICIENT * u_downwash, 0.0, 1.0)
+    # Eq. (5) in Su et al.: each motor loses a fraction b_v * U_D of its current thrust
+    loss_fraction = THRUST_DECAY_COEFFICIENT * u_downwash
     rotor_vel = data.states.rotor_vel
     k0, k1, k2 = (
         data.params.rpm2thrust[..., 0],
@@ -109,48 +109,117 @@ def downwash_fn(data: SimData) -> SimData:
     )
     return data.replace(states=states)
 
+def plot_hover_velocity_field(source_positions: np.ndarray, mass: float) -> None:
+    """Plot the far-field downwash-speed magnitude in the y=0 plane."""
+    import matplotlib.pyplot as plt
+
+    x = np.linspace(-0.6, 0.6, 300)
+    z = np.linspace(0.0, 1.15, 300)
+    X, Z = np.meshgrid(x, z)
+
+    # Every grid point lies in the y=0 plane.
+    points = np.stack((X, np.zeros_like(X), Z), axis=-1)
+    u_downwash = np.zeros_like(X)
+
+    u_hover = np.sqrt(
+        mass * GRAVITY
+        / (2.0 * AIR_DENSITY * np.pi * PROPELLER_RADIUS**2 * N_PROPELLERS)
+    )
+
+    for source_pos in source_positions:
+        source_to_point = source_pos - points
+        s = source_to_point[..., 2]
+        r = np.linalg.vector_norm(source_to_point[..., :2], axis=-1)
+        s_normalized = s / MOTOR_DISTANCE
+
+        q = np.maximum(s_normalized - S0, 1e-6)
+        half_width = S * q
+        centerline_velocity = u_hover * BD / q
+        xi = (r / MOTOR_DISTANCE) / half_width
+
+        velocity = centerline_velocity / (1.0 + (np.sqrt(2.0) - 1.0) * xi**2) ** 2
+        u_downwash += velocity
+
+    fig, ax = plt.subplots()
+    image = ax.pcolormesh(X, Z, u_downwash, shading="auto", cmap="viridis")
+    ax.scatter(source_positions[:, 0], source_positions[:, 2], color="red", label="source drone")
+    ax.set_xlabel("x (m)")
+    ax.set_ylabel("z (m)")
+    ax.set_title("Hovering-drone downwash speed")
+    ax.legend()
+    fig.colorbar(image, ax=ax, label="downward airspeed $U_D$ (m/s)")
+    plt.show()
+
 
 def main(plot: bool = True) -> None:
-    """Hover drone 0 while drone 1 flies straight through its downwash."""
+    """Hover drone 0 while drone 1 makes two downwash passes at different heights."""
     sim = Sim(n_drones=2, drone="cf21B_500", control="state")
 
     insert_fn_before(sim.step_pipeline, "integration", downwash_fn)
     sim.build_step_fn()
 
-    upper_pos = np.array([0.0, 0.0, 1.0])
-    lower_start = np.array([-0.5, 0.0, 0.5])
+    upper_pos = np.array([0.0, 0.0, 1.2])
+    outbound_height = 0.5
+    return_height = 0.95
+    lower_start = np.array([-0.5, 0.0, outbound_height])
+
     sim.data = sim.data.replace(
         states=sim.data.states.replace(pos=jnp.array([[upper_pos, lower_start]]))
     )
     sim.build_default_data()
 
-    duration = 6.0
-    speed = 1.0 / duration
     command = np.zeros((1, 2, 16))
-    command[..., 9:13] = [0.0, 0.0, 0.0, 1.0]  # level quaternion (xyzw)
+    command[..., 9:13] = [0.0, 0.0, 0.0, 1.0] 
     command[0, 0, :3] = upper_pos
+
+    waypoints = np.concatenate(
+        (
+            np.linspace(
+                lower_start,
+                [0.5, 0.0, outbound_height],
+                3 * sim.control_freq,
+                endpoint=False,
+            ),
+            np.linspace(
+                [0.5, 0.0, outbound_height],
+                [0.5, 0.0, return_height],
+                sim.control_freq,
+                endpoint=False,
+            ),
+            np.linspace(
+                [0.5, 0.0, return_height],
+                [-0.5, 0.0, return_height],
+                3 * sim.control_freq,
+            ),
+        )
+    )
+
     z_positions = []
     downwash_force_z = []
     downwash_pitch_torque = []
 
-    for step in range(int(duration * sim.control_freq)):
-        t = step / sim.control_freq
-        command[0, 1, :3] = [-0.5 + speed * t, 0.0, 0.5]
-        command[0, 1, 3:6] = [speed, 0.0, 0.0]
+    for position in waypoints:
+        command[0, 1, :3] = position
+        command[0, 1, 3:6] = 0.0  # Position-only setpoints.
+
         sim.state_control(command)
         sim.step(sim.freq // sim.control_freq)
+
         z_positions.append(np.asarray(sim.data.states.pos[0, :, 2]))
         downwash_force_z.append(np.asarray(sim.data.states.force[0, 1, 2]))
         downwash_pitch_torque.append(np.asarray(sim.data.states.torque[0, 1, 1]))
         sim.render()
 
     sim.close()
+
     if plot:
         import matplotlib.pyplot as plt
 
-        t = np.arange(len(z_positions)) / sim.control_freq
+        t = np.arange(len(waypoints)) / sim.control_freq
         z_positions = np.asarray(z_positions)
+
         fig, axes = plt.subplots(3, 1, sharex=True)
+
         axes[0].plot(t, z_positions[:, 0], label="upper drone")
         axes[0].plot(t, z_positions[:, 1], label="lower drone")
         axes[0].set_ylabel("z position (m)")
@@ -161,9 +230,18 @@ def main(plot: bool = True) -> None:
         axes[1].legend()
 
         axes[2].plot(t, downwash_pitch_torque, label="lower drone")
-        axes[2].set_xlabel("Time (s)")
+        axes[2].set_xlabel("time (s)")
         axes[2].set_ylabel("downwash pitch torque y (Nm)")
         axes[2].legend()
+
+        for axis in axes:
+            axis.axvline(3.0, color="black", linestyle="--", alpha=0.5)
+            axis.axvline(4.0, color="black", linestyle="--", alpha=0.5)
+
+        plot_hover_velocity_field(
+            np.asarray([upper_pos]),
+            float(sim.data.params.mass[0]),
+        )
         plt.show()
 
 

From 5375c2c5a5006bfa1c404374cba420b6d4672157 Mon Sep 17 00:00:00 2001
From: radu_workstation <radu141@gmail.com>
Date: Fri, 25 Sep 2026 11:55:49 +0200
Subject: [PATCH 04/10] Downwash accounts now for relative tilt of the drone
 and other minor changes

---
 examples/plugins/downwash.py | 111 ++++++++++++++++++-----------------
 1 file changed, 56 insertions(+), 55 deletions(-)

diff --git a/examples/plugins/downwash.py b/examples/plugins/downwash.py
index 991fb236..8361e830 100644
--- a/examples/plugins/downwash.py
+++ b/examples/plugins/downwash.py
@@ -1,8 +1,10 @@
 """Minimal far-field downwash external-wrench plugin.
 
-This models the downwash of level, hovering, identical Crazyflies using the
-far-field jet from Bauersfeld et al. (arXiv:2403.13321) and the thrust-decay
-model of Su et al. (arXiv:2207.09645).
+This models the downwash of identical Crazyflies using the far-field jet from
+[1] and the thrust-decay model of [2]. 
+
+[1] Bauersfeld et al. https://arxiv.org/abs/2403.13321
+[2] Su et al. https://arxiv.org/abs/2207.09645
 """
 
 from __future__ import annotations
@@ -13,7 +15,6 @@
 import numpy as np
 from jax.scipy.spatial.transform import Rotation as R
 
-from crazyflow.control.transform import motor_force2rotor_vel
 from crazyflow.sim import Sim
 from crazyflow.sim.pipeline import insert_fn_before
 
@@ -23,14 +24,12 @@
 # Physical parameters for the cf21B_500
 AIR_DENSITY = 1.225  # kg/m^3
 PROPELLER_RADIUS = 27.5e-3  # m
-MOTOR_DISTANCE = 0.035355 * 2  # m, distance between opposite motors
-N_PROPELLERS = 4
-GRAVITY = 9.81
+MOTOR_DISTANCE = 0.1  # m, distance between opposite motors
 
 # This must be fitted for the propeller/downwash setup.
-THRUST_DECAY_COEFFICIENT = 0.07 # s/m
+THRUST_DECAY_COEFFICIENT = 0.07  # s/m
 
-# Far-field fit in Eq. (9) of the Bauersfeld paper.
+# Far-field fit in Eq. (9) of [1]
 BD = 10.11
 S = 0.07668
 S0 = -5.817
@@ -40,35 +39,57 @@ def downwash_fn(data: SimData) -> SimData:
     """Apply downwash-induced thrust loss as a world-frame external wrench.
 
     The source flow originates at each drone centre, while the field is sampled
-    at every target rotor.
+    at every target rotor in the source's body frame.
     """
     rotation = R.from_quat(data.states.quat)
-    rotor_offsets = data.params.L * jnp.array(
-        [[1.0, -1.0, 0.0], [-1.0, -1.0, 0.0], [-1.0, 1.0, 0.0], [1.0, 1.0, 0.0]]
+    mixing_matrix = data.params.mixing_matrix
+
+    offsets = data.params.L * jnp.stack(
+        [-mixing_matrix[1], mixing_matrix[0], jnp.zeros_like(mixing_matrix[0])], 
+        axis=0
     )
-    rotor_offsets_world = jnp.swapaxes(rotation.as_matrix() @ rotor_offsets.T, -1, -2)
+    rotor_offsets_body = offsets.T
+    rotor_offsets_world = jnp.swapaxes(rotation.as_matrix() @ rotor_offsets_body.T, -1, -2)
     rotor_positions = data.states.pos[..., None, :] + rotor_offsets_world
 
     # Axis 1 indexes the source drone, axis 2 the target, and axis 3 its rotor.
     source_to_target = data.states.pos[:, :, None, None, :] - rotor_positions[:, None, :, :, :]
-    s = source_to_target[..., 2]  # Positive only for targets below a source.
-    r = jnp.linalg.vector_norm(source_to_target[..., :2], axis=-1)
+    # Rotate the source-minus-target displacement into each source's frame.
+    world_to_body = rotation.as_matrix().mT
+    # Broadcast each source rotation across all target drones and rotors.
+    source_to_target_body = (
+        world_to_body[:, :, None, None, :, :] @ source_to_target[..., None]
+    )[..., 0]
+    s = source_to_target_body[..., 2]
+    r = jnp.linalg.vector_norm(source_to_target_body[..., :2], axis=-1)
+
+    # Normalization according to [1] Eq. (8)
     s_normalized = s / MOTOR_DISTANCE
+    r_normalized = r / MOTOR_DISTANCE
 
     mass = data.params.mass[0]
+    gravity = -data.params.gravity_vec[2]
+    n_propellers = mixing_matrix.shape[-1]
+
     u_hover = jnp.sqrt(
-        mass * GRAVITY / (2.0 * AIR_DENSITY * jnp.pi * PROPELLER_RADIUS**2 * N_PROPELLERS)
-    )
-    half_width = S * (s_normalized - S0)
-    centerline_velocity = u_hover * BD / (s_normalized - S0)
-    xi = (r / MOTOR_DISTANCE) / half_width
-    u_downwash = centerline_velocity / (1.0 + (jnp.sqrt(2.0) - 1.0) * xi**2) ** 2
+        mass * gravity / (2.0 * AIR_DENSITY * jnp.pi * PROPELLER_RADIUS**2 * n_propellers)
+    )  # [1] Eq. (1)
+
+    # Keep the fit finite upstream, where its contribution is masked below.
+    axial_distance = jnp.maximum(s_normalized - S0, 1e-6)
+    half_width = S * axial_distance  # [1] Eq. (6)
+
+    centerline_velocity = u_hover * BD / axial_distance  # [1] Eq. (2)
+
+    xi = (r_normalized) / half_width  # [1] Eq. (4)
+
+    u_downwash = centerline_velocity / (1.0 + (jnp.sqrt(2.0) - 1.0) * xi**2) ** 2  # [1] Eq. (3)
 
     # This prevents "negative" downwash
     u_downwash = jnp.where(s_normalized > 0.1, u_downwash, 0.0)
     u_downwash = jnp.sum(u_downwash, axis=1)  # Sum all sources at each target rotor.
 
-    # Eq. (5) in Su et al.: each motor loses a fraction b_v * U_D of its current thrust
+    # [2] Eq. (5): each motor loses a fraction b_v * U_D of its current thrust
     loss_fraction = THRUST_DECAY_COEFFICIENT * u_downwash
     rotor_vel = data.states.rotor_vel
     k0, k1, k2 = (
@@ -79,37 +100,23 @@ def downwash_fn(data: SimData) -> SimData:
     motor_thrust = k0 + k1 * rotor_vel + k2 * rotor_vel**2
     thrust_delta = -loss_fraction * motor_thrust
 
-    # Map the per-motor force changes to a body-frame wrench, as in Eq. (7).
+    # Map the per-motor force changes to a body-frame wrench, as in [2] Eq. (7).
     total_thrust_delta = jnp.sum(thrust_delta, axis=-1)
     zeros = jnp.zeros_like(total_thrust_delta)
     force_body = jnp.stack((zeros, zeros, total_thrust_delta), axis=-1)
 
     lever = jnp.array([1.0, 1.0, 0.0])
-    torque_body = (data.params.mixing_matrix @ (thrust_delta * data.params.L)[..., None])[
+    torque_body = (mixing_matrix @ (thrust_delta * data.params.L)[..., None])[
         ..., 0
     ] * lever
 
-    # Account for the corresponding reaction-torque change about body z.
-    effective_motor_thrust = jnp.maximum(motor_thrust + thrust_delta, 0.0)
-    effective_rotor_vel = motor_force2rotor_vel(effective_motor_thrust, data.params.rpm2thrust)
-    c0, c1, c2 = (
-        data.params.rpm2torque[..., 0],
-        data.params.rpm2torque[..., 1],
-        data.params.rpm2torque[..., 2],
-    )
-    motor_torque = c0 + c1 * rotor_vel + c2 * rotor_vel**2
-    effective_motor_torque = c0 + c1 * effective_rotor_vel + c2 * effective_rotor_vel**2
-    reaction_torque_delta = effective_motor_torque - motor_torque
-    torque_body = torque_body + (data.params.mixing_matrix @ reaction_torque_delta[..., None])[
-        ..., 0
-    ] * jnp.array([0.0, 0.0, 1.0])
-
     states = data.states.replace(
         force=rotation.apply(force_body), torque=rotation.apply(torque_body)
     )
     return data.replace(states=states)
 
-def plot_hover_velocity_field(source_positions: np.ndarray, mass: float) -> None:
+
+def plot_hover_velocity_field(source_positions: np.ndarray, data: SimData) -> None:
     """Plot the far-field downwash-speed magnitude in the y=0 plane."""
     import matplotlib.pyplot as plt
 
@@ -121,9 +128,12 @@ def plot_hover_velocity_field(source_positions: np.ndarray, mass: float) -> None
     points = np.stack((X, np.zeros_like(X), Z), axis=-1)
     u_downwash = np.zeros_like(X)
 
+    gravity = -data.params.gravity_vec[2]
+    n_propellers = data.params.mixing_matrix.shape[-1]
+    mass = data.params.mass[0]
+
     u_hover = np.sqrt(
-        mass * GRAVITY
-        / (2.0 * AIR_DENSITY * np.pi * PROPELLER_RADIUS**2 * N_PROPELLERS)
+        mass * gravity / (2.0 * AIR_DENSITY * np.pi * PROPELLER_RADIUS**2 * n_propellers)
     )
 
     for source_pos in source_positions:
@@ -169,16 +179,13 @@ def main(plot: bool = True) -> None:
     sim.build_default_data()
 
     command = np.zeros((1, 2, 16))
-    command[..., 9:13] = [0.0, 0.0, 0.0, 1.0] 
+    command[..., 9:13] = R.from_euler("z", 0).as_quat()
     command[0, 0, :3] = upper_pos
 
     waypoints = np.concatenate(
         (
             np.linspace(
-                lower_start,
-                [0.5, 0.0, outbound_height],
-                3 * sim.control_freq,
-                endpoint=False,
+                lower_start, [0.5, 0.0, outbound_height], 3 * sim.control_freq, endpoint=False
             ),
             np.linspace(
                 [0.5, 0.0, outbound_height],
@@ -187,9 +194,7 @@ def main(plot: bool = True) -> None:
                 endpoint=False,
             ),
             np.linspace(
-                [0.5, 0.0, return_height],
-                [-0.5, 0.0, return_height],
-                3 * sim.control_freq,
+                [0.5, 0.0, return_height], [-0.5, 0.0, return_height], 3 * sim.control_freq
             ),
         )
     )
@@ -200,7 +205,6 @@ def main(plot: bool = True) -> None:
 
     for position in waypoints:
         command[0, 1, :3] = position
-        command[0, 1, 3:6] = 0.0  # Position-only setpoints.
 
         sim.state_control(command)
         sim.step(sim.freq // sim.control_freq)
@@ -238,10 +242,7 @@ def main(plot: bool = True) -> None:
             axis.axvline(3.0, color="black", linestyle="--", alpha=0.5)
             axis.axvline(4.0, color="black", linestyle="--", alpha=0.5)
 
-        plot_hover_velocity_field(
-            np.asarray([upper_pos]),
-            float(sim.data.params.mass[0]),
-        )
+        plot_hover_velocity_field(np.asarray([upper_pos]), sim.data)
         plt.show()
 
 

From d20fb994804f1b5ac9760e58a4f49027fcf80f70 Mon Sep 17 00:00:00 2001
From: radu_workstation <radu141@gmail.com>
Date: Fri, 25 Sep 2026 12:19:03 +0200
Subject: [PATCH 05/10] Cleanes up rotation matrix logic

---
 examples/plugins/downwash.py | 22 +++++++++-------------
 1 file changed, 9 insertions(+), 13 deletions(-)

diff --git a/examples/plugins/downwash.py b/examples/plugins/downwash.py
index 8361e830..bb32f174 100644
--- a/examples/plugins/downwash.py
+++ b/examples/plugins/downwash.py
@@ -1,7 +1,7 @@
 """Minimal far-field downwash external-wrench plugin.
 
 This models the downwash of identical Crazyflies using the far-field jet from
-[1] and the thrust-decay model of [2]. 
+[1] and the thrust-decay model of [2].
 
 [1] Bauersfeld et al. https://arxiv.org/abs/2403.13321
 [2] Su et al. https://arxiv.org/abs/2207.09645
@@ -41,24 +41,22 @@ def downwash_fn(data: SimData) -> SimData:
     The source flow originates at each drone centre, while the field is sampled
     at every target rotor in the source's body frame.
     """
-    rotation = R.from_quat(data.states.quat)
+    R_body_to_world = R.from_quat(data.states.quat)
+
     mixing_matrix = data.params.mixing_matrix
 
     offsets = data.params.L * jnp.stack(
-        [-mixing_matrix[1], mixing_matrix[0], jnp.zeros_like(mixing_matrix[0])], 
-        axis=0
+        [-mixing_matrix[1], mixing_matrix[0], jnp.zeros_like(mixing_matrix[0])], axis=0
     )
-    rotor_offsets_body = offsets.T
-    rotor_offsets_world = jnp.swapaxes(rotation.as_matrix() @ rotor_offsets_body.T, -1, -2)
+    rotor_offsets_world = (R_body_to_world.as_matrix() @ offsets).mT
     rotor_positions = data.states.pos[..., None, :] + rotor_offsets_world
 
     # Axis 1 indexes the source drone, axis 2 the target, and axis 3 its rotor.
     source_to_target = data.states.pos[:, :, None, None, :] - rotor_positions[:, None, :, :, :]
-    # Rotate the source-minus-target displacement into each source's frame.
-    world_to_body = rotation.as_matrix().mT
+
     # Broadcast each source rotation across all target drones and rotors.
     source_to_target_body = (
-        world_to_body[:, :, None, None, :, :] @ source_to_target[..., None]
+        R_body_to_world.as_matrix().mT[:, :, None, None, :, :] @ source_to_target[..., None]
     )[..., 0]
     s = source_to_target_body[..., 2]
     r = jnp.linalg.vector_norm(source_to_target_body[..., :2], axis=-1)
@@ -106,12 +104,10 @@ def downwash_fn(data: SimData) -> SimData:
     force_body = jnp.stack((zeros, zeros, total_thrust_delta), axis=-1)
 
     lever = jnp.array([1.0, 1.0, 0.0])
-    torque_body = (mixing_matrix @ (thrust_delta * data.params.L)[..., None])[
-        ..., 0
-    ] * lever
+    torque_body = (mixing_matrix @ (thrust_delta * data.params.L)[..., None])[..., 0] * lever
 
     states = data.states.replace(
-        force=rotation.apply(force_body), torque=rotation.apply(torque_body)
+        force=R_body_to_world.apply(force_body), torque=R_body_to_world.apply(torque_body)
     )
     return data.replace(states=states)
 

From 4f23bd4d21d4c619aa96ee9a18f3b88c66e52dc1 Mon Sep 17 00:00:00 2001
From: radu_workstation <radu141@gmail.com>
Date: Tue, 29 Sep 2026 11:45:13 +0200
Subject: [PATCH 06/10] Uses .apply function to apply rotation

---
 examples/plugins/downwash.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/examples/plugins/downwash.py b/examples/plugins/downwash.py
index bb32f174..88be8156 100644
--- a/examples/plugins/downwash.py
+++ b/examples/plugins/downwash.py
@@ -41,23 +41,23 @@ def downwash_fn(data: SimData) -> SimData:
     The source flow originates at each drone centre, while the field is sampled
     at every target rotor in the source's body frame.
     """
-    R_body_to_world = R.from_quat(data.states.quat)
+    R_world_to_body = R.from_quat(data.states.quat)
 
     mixing_matrix = data.params.mixing_matrix
 
     offsets = data.params.L * jnp.stack(
-        [-mixing_matrix[1], mixing_matrix[0], jnp.zeros_like(mixing_matrix[0])], axis=0
+        [-mixing_matrix[1], mixing_matrix[0], jnp.zeros_like(mixing_matrix[0])], axis=-1
     )
-    rotor_offsets_world = (R_body_to_world.as_matrix() @ offsets).mT
+    rotor_offsets_world = R.from_quat(data.states.quat[..., None, :]).apply(offsets)
     rotor_positions = data.states.pos[..., None, :] + rotor_offsets_world
 
     # Axis 1 indexes the source drone, axis 2 the target, and axis 3 its rotor.
     source_to_target = data.states.pos[:, :, None, None, :] - rotor_positions[:, None, :, :, :]
 
     # Broadcast each source rotation across all target drones and rotors.
-    source_to_target_body = (
-        R_body_to_world.as_matrix().mT[:, :, None, None, :, :] @ source_to_target[..., None]
-    )[..., 0]
+    source_to_target_body = R.from_quat(data.states.quat[..., None, None, :]).apply(
+        source_to_target, inverse=True
+    )
     s = source_to_target_body[..., 2]
     r = jnp.linalg.vector_norm(source_to_target_body[..., :2], axis=-1)
 
@@ -107,7 +107,7 @@ def downwash_fn(data: SimData) -> SimData:
     torque_body = (mixing_matrix @ (thrust_delta * data.params.L)[..., None])[..., 0] * lever
 
     states = data.states.replace(
-        force=R_body_to_world.apply(force_body), torque=R_body_to_world.apply(torque_body)
+        force=R_world_to_body.apply(force_body), torque=R_world_to_body.apply(torque_body)
     )
     return data.replace(states=states)
 

From 6d2f21af323280b9fa1d4ac9b68b219f2b56d9f0 Mon Sep 17 00:00:00 2001
From: radu_workstation <radu141@gmail.com>
Date: Tue, 29 Sep 2026 16:38:44 +0200
Subject: [PATCH 07/10] Adds scaling factor for tilt angle between source and
 target

---
 examples/plugins/downwash.py | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/examples/plugins/downwash.py b/examples/plugins/downwash.py
index 88be8156..aad6211c 100644
--- a/examples/plugins/downwash.py
+++ b/examples/plugins/downwash.py
@@ -85,6 +85,11 @@ def downwash_fn(data: SimData) -> SimData:
 
     # This prevents "negative" downwash
     u_downwash = jnp.where(s_normalized > 0.1, u_downwash, 0.0)
+
+    # Account for inclination between target drone and source field
+    z_axes = R_world_to_body.as_matrix()[..., :, 2]
+    cos_theta = jnp.sum(z_axes[:, :, None, :] * z_axes[:, None, :, :], axis=-1)
+    u_downwash = u_downwash * cos_theta[..., None]
     u_downwash = jnp.sum(u_downwash, axis=1)  # Sum all sources at each target rotor.
 
     # [2] Eq. (5): each motor loses a fraction b_v * U_D of its current thrust
@@ -100,6 +105,11 @@ def downwash_fn(data: SimData) -> SimData:
 
     # Map the per-motor force changes to a body-frame wrench, as in [2] Eq. (7).
     total_thrust_delta = jnp.sum(thrust_delta, axis=-1)
+
+    # Account for inclination between target drone and source field
+    z_axes = R_world_to_body.as_matrix()[..., :, 2]
+    cos_theta = jnp.sum(z_axes[:, :, None, :] * z_axes[:, None, :, :], axis=-1)
+
     zeros = jnp.zeros_like(total_thrust_delta)
     force_body = jnp.stack((zeros, zeros, total_thrust_delta), axis=-1)
 
@@ -181,17 +191,18 @@ def main(plot: bool = True) -> None:
     waypoints = np.concatenate(
         (
             np.linspace(
-                lower_start, [0.5, 0.0, outbound_height], 3 * sim.control_freq, endpoint=False
+                lower_start, [0.3, 0.0, outbound_height], 3 * sim.control_freq, endpoint=False
             ),
+            np.tile([0.3, 0.0, outbound_height], (2 * sim.control_freq, 1)),
+            np.linspace([0.3, 0.0, outbound_height], lower_start, int(0.5 * sim.control_freq)),
+            np.tile(lower_start, (2 * sim.control_freq, 1)),
+            np.linspace(lower_start, [-0.5, 0.0, return_height], sim.control_freq, endpoint=False),
             np.linspace(
-                [0.5, 0.0, outbound_height],
+                [-0.5, 0.0, return_height],
                 [0.5, 0.0, return_height],
-                sim.control_freq,
+                3 * sim.control_freq,
                 endpoint=False,
             ),
-            np.linspace(
-                [0.5, 0.0, return_height], [-0.5, 0.0, return_height], 3 * sim.control_freq
-            ),
         )
     )
 

From bd12101ff458e6dcb3ab5931b07e653d07345c38 Mon Sep 17 00:00:00 2001
From: radu_workstation <radu141@gmail.com>
Date: Tue, 29 Sep 2026 17:08:40 +0200
Subject: [PATCH 08/10] Adds drag model

---
 examples/plugins/downwash.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/examples/plugins/downwash.py b/examples/plugins/downwash.py
index aad6211c..ce8bfae3 100644
--- a/examples/plugins/downwash.py
+++ b/examples/plugins/downwash.py
@@ -44,6 +44,7 @@ def downwash_fn(data: SimData) -> SimData:
     R_world_to_body = R.from_quat(data.states.quat)
 
     mixing_matrix = data.params.mixing_matrix
+    drag_matrix = data.params.drag_matrix
 
     offsets = data.params.L * jnp.stack(
         [-mixing_matrix[1], mixing_matrix[0], jnp.zeros_like(mixing_matrix[0])], axis=-1
@@ -89,6 +90,14 @@ def downwash_fn(data: SimData) -> SimData:
     # Account for inclination between target drone and source field
     z_axes = R_world_to_body.as_matrix()[..., :, 2]
     cos_theta = jnp.sum(z_axes[:, :, None, :] * z_axes[:, None, :, :], axis=-1)
+
+    # Wind at each target rotor summed over sources
+    rotor_wind_world = jnp.sum(-u_downwash[..., None] * z_axes[:, :, None, None, :], axis=1)
+
+    # Rotor-averaged wind representing the flow at the target CoM.
+    wind_com_world = jnp.mean(rotor_wind_world, axis=2)
+    wind_com_body = R_world_to_body.apply(wind_com_world, inverse=True)
+
     u_downwash = u_downwash * cos_theta[..., None]
     u_downwash = jnp.sum(u_downwash, axis=1)  # Sum all sources at each target rotor.
 
@@ -113,6 +122,10 @@ def downwash_fn(data: SimData) -> SimData:
     zeros = jnp.zeros_like(total_thrust_delta)
     force_body = jnp.stack((zeros, zeros, total_thrust_delta), axis=-1)
 
+    # Compute drag induced through downwash
+    drag_body = (-drag_matrix @ wind_com_body[..., None])[..., 0]
+    force_body += drag_body
+
     lever = jnp.array([1.0, 1.0, 0.0])
     torque_body = (mixing_matrix @ (thrust_delta * data.params.L)[..., None])[..., 0] * lever
 

From e09fd5355420399de58cbcc9ff87fda11d3baa94 Mon Sep 17 00:00:00 2001
From: radu_workstation <radu141@gmail.com>
Date: Tue, 29 Sep 2026 17:22:36 +0200
Subject: [PATCH 09/10] Removes lines from plotting

---
 examples/plugins/downwash.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/examples/plugins/downwash.py b/examples/plugins/downwash.py
index ce8bfae3..e0e0b1df 100644
--- a/examples/plugins/downwash.py
+++ b/examples/plugins/downwash.py
@@ -258,10 +258,6 @@ def main(plot: bool = True) -> None:
         axes[2].set_ylabel("downwash pitch torque y (Nm)")
         axes[2].legend()
 
-        for axis in axes:
-            axis.axvline(3.0, color="black", linestyle="--", alpha=0.5)
-            axis.axvline(4.0, color="black", linestyle="--", alpha=0.5)
-
         plot_hover_velocity_field(np.asarray([upper_pos]), sim.data)
         plt.show()
 

From a35ca03f8f7865bf95248926b0d0929ccc7a1eb3 Mon Sep 17 00:00:00 2001
From: radu_workstation <radu141@gmail.com>
Date: Tue, 29 Sep 2026 17:36:30 +0200
Subject: [PATCH 10/10] Adds an additional sample point for the downwash field
 at the drone COM for accurate computation of the local wind velocity vector

---
 examples/plugins/downwash.py | 26 +++++++++++---------------
 1 file changed, 11 insertions(+), 15 deletions(-)

diff --git a/examples/plugins/downwash.py b/examples/plugins/downwash.py
index e0e0b1df..0f82df2c 100644
--- a/examples/plugins/downwash.py
+++ b/examples/plugins/downwash.py
@@ -39,7 +39,7 @@ def downwash_fn(data: SimData) -> SimData:
     """Apply downwash-induced thrust loss as a world-frame external wrench.
 
     The source flow originates at each drone centre, while the field is sampled
-    at every target rotor in the source's body frame.
+    at every target rotor and CoM in the source's body frame.
     """
     R_world_to_body = R.from_quat(data.states.quat)
 
@@ -52,10 +52,12 @@ def downwash_fn(data: SimData) -> SimData:
     rotor_offsets_world = R.from_quat(data.states.quat[..., None, :]).apply(offsets)
     rotor_positions = data.states.pos[..., None, :] + rotor_offsets_world
 
-    # Axis 1 indexes the source drone, axis 2 the target, and axis 3 its rotor.
-    source_to_target = data.states.pos[:, :, None, None, :] - rotor_positions[:, None, :, :, :]
+    sample_positions = jnp.concatenate([rotor_positions, data.states.pos[..., None, :]], axis=2)
 
-    # Broadcast each source rotation across all target drones and rotors.
+    # Axis 1 indexes the source, axis 2 the target, and axis 3 its rotors then CoM.
+    source_to_target = data.states.pos[:, :, None, None, :] - sample_positions[:, None, :, :, :]
+
+    # Broadcast each source rotation across all target drones and sampling points.
     source_to_target_body = R.from_quat(data.states.quat[..., None, None, :]).apply(
         source_to_target, inverse=True
     )
@@ -87,17 +89,15 @@ def downwash_fn(data: SimData) -> SimData:
     # This prevents "negative" downwash
     u_downwash = jnp.where(s_normalized > 0.1, u_downwash, 0.0)
 
-    # Account for inclination between target drone and source field
     z_axes = R_world_to_body.as_matrix()[..., :, 2]
-    cos_theta = jnp.sum(z_axes[:, :, None, :] * z_axes[:, None, :, :], axis=-1)
 
-    # Wind at each target rotor summed over sources
-    rotor_wind_world = jnp.sum(-u_downwash[..., None] * z_axes[:, :, None, None, :], axis=1)
-
-    # Rotor-averaged wind representing the flow at the target CoM.
-    wind_com_world = jnp.mean(rotor_wind_world, axis=2)
+    # The final sample is the CoM; each source's wind follows its negative z-axis.
+    wind_com_world = jnp.sum(-u_downwash[..., -1, None] * z_axes[:, :, None, :], axis=1)
     wind_com_body = R_world_to_body.apply(wind_com_world, inverse=True)
 
+    # Project only the rotor samples onto the target axis for thrust loss.
+    u_downwash = u_downwash[..., :-1]
+    cos_theta = jnp.sum(z_axes[:, :, None, :] * z_axes[:, None, :, :], axis=-1)
     u_downwash = u_downwash * cos_theta[..., None]
     u_downwash = jnp.sum(u_downwash, axis=1)  # Sum all sources at each target rotor.
 
@@ -115,10 +115,6 @@ def downwash_fn(data: SimData) -> SimData:
     # Map the per-motor force changes to a body-frame wrench, as in [2] Eq. (7).
     total_thrust_delta = jnp.sum(thrust_delta, axis=-1)
 
-    # Account for inclination between target drone and source field
-    z_axes = R_world_to_body.as_matrix()[..., :, 2]
-    cos_theta = jnp.sum(z_axes[:, :, None, :] * z_axes[:, None, :, :], axis=-1)
-
     zeros = jnp.zeros_like(total_thrust_delta)
     force_body = jnp.stack((zeros, zeros, total_thrust_delta), axis=-1)