estimator_diff

A differentiation-based filter.

This estimator is originally from Chris McKinnon. It was adapted to work with similarly to the new estimators for backwards compatability.

The filter itself is the pure function diff_correct, which the DiffFilter class jits. Quaternions are scalar last and all rotations use scipy's Rotation.

DiffData

State of the differentiation-based filter.

create_empty(batch_shape=()) classmethod

Create the initial state.

Source code in drone_estimators/estimator_diff.py

@classmethod
def create_empty(cls, batch_shape: tuple[int, ...] = ()) -> DiffData:
    """Create the initial state."""
    zeros = np.zeros((*batch_shape, 3))
    quat = np.tile([0.0, 0.0, 0.0, 1.0], (*batch_shape, 1))
    return cls(
        pos=zeros,
        vel=zeros,
        acc=zeros,
        quat=quat,
        ang_vel=zeros,
        vel_meas=zeros,
        acc_meas=zeros,
        ang_vel_meas=zeros,
        pos_prev=zeros,
        vel_prev=zeros,
        quat_prev=quat,
    )

DiffFilter(filter_parameters, batch_shape=(), jit=True)

Bases: object

Vicon state estimation and filtering.

Parameters

filter_parameters : sequence of floats The 5 time constants of the low pass filters, see DiffSettings batch_shape : tuple of ints The batch dimensions, e.g., (n_drones,). All arrays have leading batch dimensions. jit : bool If the filter should be jit compiled with jax.

Create the filter and compile the correction step if jit is enabled.

Source code in drone_estimators/estimator_diff.py

def __init__(
    self, filter_parameters: tuple, batch_shape: tuple[int, ...] = (), jit: bool = True
):
    """Create the filter and compile the correction step if jit is enabled."""
    # The numeric derivatives divide by dt (twice for the acceleration). In 32 bit, this
    # amplifies the rounding errors of the measurements too much.
    if jit and not jax.config.jax_enable_x64:
        raise RuntimeError(
            "The differentiation-based filter needs 64 bit precision. "
            'Call jax.config.update("jax_enable_x64", True) before creating it.'
        )
    self.settings = DiffSettings(*filter_parameters)
    self.data = DiffData.create_empty(batch_shape)
    self.dt = np.zeros(batch_shape)  # Time since the last correction

    self._correct = diff_correct
    if jit:
        self.data = jax.tree.map(jnp.asarray, self.data)
        self._correct = jax.jit(diff_correct)
        # Compile now instead of when the first measurement arrives
        pos, quat, mask = self.data.pos, self.data.quat, np.zeros(batch_shape, dtype=bool)
        self._correct(self.data, self.settings, pos, quat, self.dt, mask)
    self._update_estimate()

correct(pos, quat, mask=None)

This function is not part of the original filter and only for compatability.

If given, only batch elements where mask is True are corrected.

Source code in drone_estimators/estimator_diff.py

def correct(self, pos: Array, quat: Array, mask: Array | None = None) -> EstimatorData:
    """This function is not part of the original filter and only for compatability.

    If given, only batch elements where mask is True are corrected.
    """
    pos = np.asarray(pos, dtype=np.float64)
    quat = np.asarray(quat, dtype=np.float64)
    self.data = self._correct(self.data, self.settings, pos, quat, self.dt, mask)
    self.dt = np.where(True if mask is None else mask, 0.0, self.dt)
    self._update_estimate()
    return self._estimate

predict(dt)

This function is not part of the original filter and only for compatability.

Source code in drone_estimators/estimator_diff.py

def predict(self, dt: float) -> EstimatorData:
    """This function is not part of the original filter and only for compatability."""
    # Since the original filter doesn't inherently support the prediction/correction
    # form of a Kalman filter, we only step the time in the prediction step. In the
    # correction step, we use the accumulated time step size to correct with the
    # correct dt. Therefore, the estimate itself only changes in the correction step.
    self.dt += dt
    return self._estimate

set_state(pos, quat)

This function is not part of the original filter and only for compatability.

Source code in drone_estimators/estimator_diff.py

def set_state(self, pos: Array, quat: Array):
    """This function is not part of the original filter and only for compatability."""
    pos = np.asarray(pos, dtype=np.float64)
    quat = np.asarray(quat, dtype=np.float64)
    # Also use the state as the previous measurement. Otherwise, the first numeric
    # derivatives are computed from zero and cause a huge spike in the velocities.
    self.data = self.data.replace(pos=pos, quat=quat, pos_prev=pos, quat_prev=quat)
    self._update_estimate()

DiffSettings

Time constants of the low pass filters.

compute_numerical_derivatives(data, pos, quat, dt)

Compute velocity, acceleration, and angular velocity from consecutive measurements.

Source code in drone_estimators/estimator_diff.py

def compute_numerical_derivatives(data: DiffData, pos: Array, quat: Array, dt: float) -> DiffData:
    """Compute velocity, acceleration, and angular velocity from consecutive measurements."""
    xp = data.pos.__array_namespace__()

    # Skip impossibly small time differences. To stay jittable, we compute the derivatives with a
    # dummy dt and keep the previous values instead of returning early.
    valid = dt > 1e-15
    dt = xp.where(valid, dt, 1.0)
    vel_meas = (pos - data.pos_prev) / dt
    acc_meas = (vel_meas - data.vel_prev) / dt
    ang_vel_meas = (R.from_quat(data.quat_prev).inv() * R.from_quat(quat)).as_rotvec() / dt

    return data.replace(
        vel_meas=xp.where(valid, vel_meas, data.vel_meas),
        acc_meas=xp.where(valid, acc_meas, data.acc_meas),
        ang_vel_meas=xp.where(valid, ang_vel_meas, data.ang_vel_meas),
        pos_prev=xp.where(valid, pos, data.pos_prev),
        vel_prev=xp.where(valid, vel_meas, data.vel_prev),
        quat_prev=xp.where(valid, quat, data.quat_prev),
    )

diff_correct(data, settings, pos, quat, dt, mask=None)

Correct the estimate with a new measurement.

Parameters:

Name	Type	Description	Default
data	DiffData	The current state of the estimator.	required
settings	DiffSettings	The filter time constants.	required
pos	Array	The measured position.	required
quat	Array	The measured orientation.	required
dt	float	The time since the last correction.	required
mask	Array | None	Optional, only batch elements where mask is True are corrected.	None

Returns:

Type	Description
DiffData	The corrected state.

Source code in drone_estimators/estimator_diff.py

def diff_correct(
    data: DiffData,
    settings: DiffSettings,
    pos: Array,
    quat: Array,
    dt: float,
    mask: Array | None = None,
) -> DiffData:
    """Correct the estimate with a new measurement.

    Args:
        data: The current state of the estimator.
        settings: The filter time constants.
        pos: The measured position.
        quat: The measured orientation.
        dt: The time since the last correction.
        mask: Optional, only batch elements where mask is True are corrected.

    Returns:
        The corrected state.
    """
    xp = data.pos.__array_namespace__()
    data_in = data
    if mask is not None:
        # Batch elements without a measurement might contain invalid values, e.g., zero quaternions
        pos = xp.where(mask[..., None], pos, data.pos)
        quat = xp.where(mask[..., None], quat, data.quat)
    dt = xp.asarray(dt)[..., None]  # Broadcast against the last axis
    data = compute_numerical_derivatives(data, pos, quat, dt)
    data = prior_update(data, dt)
    data = low_pass_filter(data, settings, pos, quat, dt)
    return data if mask is None else select(mask, data, data_in)

low_pass_filter(data, settings, pos, quat, dt)

Move the predicted state towards the measurements and their numerical derivatives.

Each quantity moves by the fraction 1 - exp(-dt / tau) of its difference to the measurement.

Source code in drone_estimators/estimator_diff.py

def low_pass_filter(
    data: DiffData, settings: DiffSettings, pos: Array, quat: Array, dt: float
) -> DiffData:
    """Move the predicted state towards the measurements and their numerical derivatives.

    Each quantity moves by the fraction 1 - exp(-dt / tau) of its difference to the measurement.
    """
    xp = data.pos.__array_namespace__()

    def weight(tau: float) -> Array:
        return 1.0 - xp.exp(-dt / tau)

    pos = data.pos + weight(settings.tau_est_trans) * (pos - data.pos)
    vel = data.vel + weight(settings.tau_est_trans_dot) * (data.vel_meas - data.vel)
    acc = data.acc + weight(settings.tau_est_trans_dot_dot) * (data.acc_meas - data.acc)
    ang_vel = data.ang_vel + weight(settings.tau_est_rot_dot) * (data.ang_vel_meas - data.ang_vel)
    # The orientation moves along the rotation from the predicted to the measured one
    rot = R.from_quat(data.quat)
    rotvec = weight(settings.tau_est_rot) * (rot.inv() * R.from_quat(quat)).as_rotvec()
    quat = (rot * R.from_rotvec(rotvec)).as_quat()
    return data.replace(pos=pos, vel=vel, acc=acc, quat=quat, ang_vel=ang_vel)

prior_update(data, dt)

Predict the state assuming constant acceleration and angular velocity.

Source code in drone_estimators/estimator_diff.py

def prior_update(data: DiffData, dt: float) -> DiffData:
    """Predict the state assuming constant acceleration and angular velocity."""
    pos = data.pos + dt * data.vel + 0.5 * dt * dt * data.acc
    vel = data.vel + dt * data.acc
    quat = (R.from_quat(data.quat) * R.from_rotvec(data.ang_vel * dt)).as_quat()
    return data.replace(pos=pos, vel=vel, quat=quat)