Control allocation w/ smoothed saturation

def _pd_controller(y, sp=None, mass=None):
    """Cascaded PD: pos error → desired attitude → torques. Returns (T, tau_x, tau_y, tau_z).

    BA-2 (2026-04-29): added optional ``mass`` kwarg so the
    factory variants in ``make_rhs_A3`` / ``make_rhs_A7`` can override
    the module-global ``MASS`` for the hover-thrust feed-forward term.
    THRUST_MAX is held at the parent airframe's value (39.24 N) since it
    represents the physical thrust ceiling; airframe-mass perturbations
    in {0.8..1.2} kg stay well within this envelope.
    """
    if sp is None:
        sp = HOVER_SP
    eff_mass = MASS if mass is None else mass
    px, py, pz = y[0], y[1], y[2]
    vx, vy, vz = y[3], y[4], y[5]
    phi, theta, psi = y[6], y[7], y[8]
    p, q, r = y[9], y[10], y[11]

    ax_d = KP_POS * (sp[0] - px) + KD_POS * (sp[3] - vx)
    ay_d = KP_POS * (sp[1] - py) + KD_POS * (sp[4] - vy)
    az_d = KP_POS * (sp[2] - pz) + KD_POS * (sp[5] - vz)

    T_des = eff_mass * (G + az_d)
    phi_des = (1.0 / G) * (ax_d * np.sin(psi) - ay_d * np.cos(psi))
    theta_des = (1.0 / G) * (ax_d * np.cos(psi) + ay_d * np.sin(psi))

    tau_x = KP_ATT * np.arctan2(np.sin(phi_des - phi), np.cos(phi_des - phi)) - KD_ATT * p
    tau_y = KP_ATT * np.arctan2(np.sin(theta_des - theta), np.cos(theta_des - theta)) - KD_ATT * q
    tau_z = KP_YAW * np.arctan2(np.sin(-psi), np.cos(-psi)) - KD_YAW * r

    T_des = np.clip(T_des, 0.0, THRUST_MAX)
    tau_x = np.clip(tau_x, -TORQUE_CLIP, TORQUE_CLIP)
    tau_y = np.clip(tau_y, -TORQUE_CLIP, TORQUE_CLIP)
    tau_z = np.clip(tau_z, -TORQUE_CLIP * 0.25, TORQUE_CLIP * 0.25)
    return T_des, tau_x, tau_y, tau_z

def _body_forces(T, phi, theta, psi):
    """Thrust-to-inertial force components."""
    cp, sp = np.cos(phi), np.sin(phi)
    ct, st = np.cos(theta), np.sin(theta)
    cy, sy = np.cos(psi), np.sin(psi)
    Fx = T * (cy * st * cp + sy * sp)
    Fy = T * (sy * st * cp - cy * sp)
    Fz = T * ct * cp
    return Fx, Fy, Fz

def _euler_kinematics(phi, theta, p, q, r):
    """Euler-angle rates from body rates. Returns (dphi, dtheta, dpsi)."""
    cp, sp = np.cos(phi), np.sin(phi)
    theta_c = np.clip(theta, -1.39, 1.39)
    tan_th = np.tan(theta_c)
    cos_th = np.cos(theta_c)
    sec_th = 1.0 / cos_th if abs(cos_th) > 1e-12 else 1e12 * np.sign(cos_th)
    dphi = p + q * sp * tan_th + r * cp * tan_th
    dtheta = q * cp - r * sp
    dpsi = (q * sp + r * cp) * sec_th
    return dphi, dtheta, dpsi

def _quad12(y, T, tau_x, tau_y, tau_z, mass=None):
    """Core 12-state quadrotor dynamics. Returns d[0:12].

    BA-2 (2026-04-29): added optional ``mass`` kwarg so the
    factory variants can override the module-global ``MASS`` for the
    translational acceleration / drag terms.
    """
    eff_mass = MASS if mass is None else mass
    phi, theta, psi = y[6], y[7], y[8]
    p, q, r = y[9], y[10], y[11]
    Fx, Fy, Fz = _body_forces(T, phi, theta, psi)

    d = np.empty(12)
    d[0] = y[3]; d[1] = y[4]; d[2] = y[5]
    d[3] = (Fx - CD * y[3]) / eff_mass
    d[4] = (Fy - CD * y[4]) / eff_mass
    d[5] = (Fz - CD * y[5]) / eff_mass - G
    d[6], d[7], d[8] = _euler_kinematics(phi, theta, p, q, r)
    d[9] = (tau_x + (IYY - IZZ) * q * r) / IXX
    d[10] = (tau_y + (IZZ - IXX) * p * r) / IYY
    d[11] = (tau_z + (IXX - IYY) * p * q) / IZZ
    return d

def _smooth_sat(x, lim):
    return lim * np.tanh(x / lim)

def rhs_B5(t, y):
    body = y[:12]
    alloc = y[12:18]

    T, tx, ty, tz = _pd_controller(body)
    cmd = np.array([T, tx, ty, tz, tx * 0.5, ty * 0.5])
    saturated = _smooth_sat(cmd, _ALLOC_SAT)
    d_alloc = (saturated - alloc) / _ALLOC_TAU

    T_eff, tx_eff, ty_eff, tz_eff = alloc[0], alloc[1], alloc[2], alloc[3]
    d_body = _quad12(body, T_eff, tx_eff, ty_eff, tz_eff)

    return np.concatenate([d_body, d_alloc])