import numpy as np
from scipy.spatial.transform import Rotation as R, Slerp
import matplotlib.pyplot as plt
from scipy.linalg import block_diag

from dataclasses import dataclass
import copy

from ipynb.fs.defs.data_simulation import generate_swing_path, generate_imu_measurements, G_GLOBAL
from ipynb.fs.defs.plot_estimated_data import plot_eskf_results, plot_velocity_kinematics, plot_g_error, plot_gyro_bias_error

ESKF

T=2.0
fs = 5000 # frequency of truth-data
truth_data = generate_swing_path(T, fs)

target_fs=100 # noisy IMU data frequency
imu_data = generate_imu_measurements(
    truth_data, 
    target_fs, 
    acc_noise=0.1,    # m/s^2
    gyro_noise=0.01   # rad/s
)
truth_data['acc_global'].shape, imu_data['acc'].shape

((13709, 3), (280, 3))

ESKF Variable and Unit Specification (SI Standards)

Category	Variable	Symbol	Unit (SI)	Type	Frame
Nominal State	Position	$p$	meters ($m$)	Vector3	World
	Velocity	$v$	$m/s$	Vector3	World
	Orientation	$q$	unitless	Quaternion (x,y,z,w)	Body $\to$ World
	Accel Bias	$a_b$	$m/s^2$	Vector3	Body
	Gyro Bias	$\omega_b$	$rad/s$	Vector3	Body
	Gravity	$g$	$m/s^2$	Vector3	World
Error State	Pos. Error	$\delta p$	$m$	Vector3	World
	Vel. Error	$\delta v$	$m/s$	Vector3	World
	Angular Error	$\delta \theta$	radians ($rad$)	Vector3	Body
	A-Bias Error	$\delta a_b$	$m/s^2$	Vector3	Body
	G-Bias Error	$\delta \omega_b$	$rad/s$	Vector3	Body
	Gravity Error	$\delta g$	$m/s^2$	Vector3	World
Measurements	Raw Accel	$a_{imu}$	$m/s^2$	Vector3	Body
	Raw Gyro	$\omega_{imu}$	$rad/s$	Vector3	Body
	Time Step	$\Delta t$	seconds ($s$)	Scalar	N/A
Noise Params	Accel Noise	$\sigma_a$	$m/s^2$	Std. Dev	Body
	Gyro Noise	$\sigma_\omega$	$rad/s$	Std. Dev	Body
	Accel Walk	$\sigma_{ab}$	$m/s^2 \sqrt{s}$	Random Walk	Body
	Gyro Walk	$\sigma_{\omega b}$	$rad/s \sqrt{s}$	Random Walk	Body

Critical Notes for Implementation:

1. Angular Error ($\delta \theta$): In the ESKF, we do not use a 4D quaternion error. We use a 3D rotation vector (angle-axis) representing a small local rotation.
2. Body Frame: Accelerometer and Gyroscope measurements (and their biases) are strictly in the sensor's local frame.
3. Gravity: Since $g$ is in the state, its value is initialized with the first measurements.

@dataclass
class NominalState:
    p: np.ndarray       # Position (3,)
    v: np.ndarray       # Velocity (3,)
    q: R                # Orientation (SciPy Rotation object)
    acc_bias: np.ndarray # (3,)
    gyro_bias: np.ndarray # (3,)
    g: np.ndarray       # Gravity vector (3,)
    
    def to_array(self):
        # Useful for logging: [p, v, q_quat, ab, wb, g]
        return np.concatenate([self.p, self.v, self.q.as_quat(), 
                               self.acc_bias, self.gyro_bias, self.g])

@dataclass
class IMUInput:
    acc_raw: np.ndarray  # Raw specific force [m/s^2]
    gyro_raw: np.ndarray # Raw angular velocity [rad/s]
    dt: float            # Time step [s]

ESKF

The State Vector Design:

• Nominal State ($x$): $p, v, q, a_b, \omega_b, g \in \mathbb{R}^{19}$
• Error State ($\delta x$): $[\delta p, \delta v, \delta \theta, \delta a_b, \delta \omega_b, \delta g]^T \in \mathbb{R}^{18}$

ESKF Error Transition Matrix ($F_x$) - Block Documentation

The 18x18 Jacobian matrix $F_x$ defines the discrete-time evolution of the error state:
$\delta x_{k+1} \approx F_x \cdot \delta x_k$

State Block	Dependency	Mathematical Term	Physical Interpretation
Position ($\delta p$)	Velocity	$I \cdot \Delta t$	Velocity error integrates linearly into position error.
Velocity ($\delta v$)	Orientation	$-R[a_u]_\times \Delta t$	Tilt-Gravity Coupling: If orientation is wrong, gravity is subtracted along the wrong axis, creating a massive velocity drift.
	Accel Bias	$-R \cdot \Delta t$	Uncorrected sensor bias integrates directly into velocity.
	Gravity	$I \cdot \Delta t$	Errors in the estimated gravity vector translate directly to velocity.
Orientation ($\delta \theta$)	Orientation	$Exp(\omega_u \Delta t)^T$	The current orientation error is rotated by the racket's own motion (the "gyro sweep").
	Gyro Bias	$-I \cdot \Delta t$	Uncorrected gyro bias integrates into a growing angle error.
Accel Bias ($\delta a_b$)	Accel Bias	$I$	Modeled as a Random Walk (constant in the Jacobian).
Gyro Bias ($\delta \omega_b$)	Gyro Bias	$I$	Modeled as a Random Walk (constant in the Jacobian).
Gravity ($\delta g$)	Gravity	$I$	Modeled as a constant/slow-drifting vector in the world frame.

Implementation Note: All terms are multiplied by $\Delta t$ because we are converting a continuous-time derivative into a discrete-time step.

ESKF Process Noise Matrix ($Q$) - Parameter Mapping

The matrix $Q$ represents the covariance of the noise injected into the error state over a time step $\Delta t$. We assume the noise sources are uncorrelated, resulting in a diagonal-block matrix.

State Affected	Parameter	Units	Discrete-Time Variance ($Q$ diagonal)	Physical Source
Velocity ($\delta v$)	$\sigma_a$	$m/s^2$	$(\sigma_a \cdot \Delta t)^2$	Accelerometer "White Noise"
Orientation ($\delta \theta$)	$\sigma_\omega$	$rad/s$	$(\sigma_\omega \cdot \Delta t)^2$	Gyroscope "White Noise"
Accel Bias ($\delta a_b$)	$\sigma_{ab}$	$m/s^2\sqrt{s}$	$\sigma_{ab}^2 \cdot \Delta t$	Accel Bias "Random Walk"
Gyro Bias ($\delta \omega_b$)	$\sigma_{\omega b}$	$rad/s\sqrt{s}$	$\sigma_{\omega b}^2 \cdot \Delta t$	Gyro Bias "Random Walk"
Gravity ($\delta g$)	$\sigma_g$	$m/s^2$	$\sigma_g^2 \cdot \Delta t$	Gravity instability (usually $\approx 0$)

Note on Units: - $\sigma_a$ and $\sigma_\omega$ are the standard deviations of the raw sensor noise.

• $\sigma_{ab}$ and $\sigma_{\omega b}$ are the "Random Walk" densities. If the bias is expected to drift by $0.1$ units over $100$ seconds, $\sigma \approx 0.1 / \sqrt{100}$.

Why we distinguish between $dt$ and $dt^2$:

• White Noise ($dt^2$): For velocity and orientation, the noise is added to the derivative. When you integrate the derivative over $dt$, the variance of the resulting state grows by $dt^2$ (since $Var(k \cdot X) = k^2 Var(X)$).
• Random Walk ($dt$): For biases, we model the change itself as a Wiener process. The variance of a Wiener process grows linearly with time ($dt$).

Pseudo "Still measurement"

The "Unobservable" Trap

Because Tilt is locked $(P_{6:9} \approx 0)$, and both Bias and Gravity are being estimated from the same 3-axis accelerometer, there is a mathematical ambiguity. The filter might struggle to know if a $0.05 \text{ m/s}^2$ error is a "Gravity vector slightly off" or "Bias is slightly off."The Result: You might see the Gravity vector and the Bias vector both move slightly.This is actually fine for short-term tracking! As long as the sum of (Rotated Gravity + Bias) equals the static IMU reading, your velocity will stay at zero.

Nominal State Kinematics for the Prediction Step

The nominal state kinematics is implemented in the integrator strategy.

The Physics Model

Derivatives

The continuous-time dynamics of the nominal state are computed via `kinematics.get_derivatives. This defines the "unbiased" physics of the system: $$\dot{\mathbf{p}} = \mathbf{v}$$$$\dot{\mathbf{v}} = \mathbf{R}_t (\mathbf{a}_m - \mathbf{a}_b) + \mathbf{g}$$$$\boldsymbol{\omega}_u = \boldsymbol{\omega}_m - \boldsymbol{\omega}_b$$

• $\mathbf{a}_m, \boldsymbol{\omega}_m$: Raw measurements from the IMU.
• $\mathbf{a}_b, \boldsymbol{\omega}_b$: The current estimates of sensor biases.
• $\mathbf{R}_t$: The Rotation Matrix representing the orientation from the Body frame to the World frame.

Integration for Nominal Kinematics

Numerical integration is handled by the `DiscreteRK4Integrator. This high-order approach provides a numerical solution to the differential equations governing IMU physics, capturing the "arc" of high-speed motions (like a tennis swing) that simple Euler integration would miss.

The RK4 Methodology

The integrator samples the state derivatives at the start, midpoints, and end of the interval $\Delta t$ to compute four distinct slopes ($k_1$ through $k_4$).For position and velocity, the state is updated via a weighted average of these slopes:$$\mathbf{x}_{t+\Delta t} = \mathbf{x}_t + \frac{\Delta t}{6}(k_1 + 2k_2 + 2k_3 + k_4)$$

Velocity and Rotation Coupling

While the unbiased angular velocity $\boldsymbol{\omega}_u$ and acceleration $\mathbf{a}_u$ are constant over the discrete interval $\Delta t$ (Zero-Order Hold), the World-frame acceleration is time-varying because the body is rotating.

Note, that const. means here: $\boldsymbol{\omega}_i = \boldsymbol{\omega}_1 = \boldsymbol{\omega}_2= \boldsymbol{\omega}_3 = \boldsymbol{\omega}_4$

The velocity derivatives $\dot{\mathbf{v}}$ are computed using distinct rotation matrices $\mathbf{R}(\mathbf{q}_1) \dots \mathbf{R}(\mathbf{q}_4)$. These correspond to the nominal orientations sampled throughout the interval:

• $\mathbf{q}_1$: Orientation at $t$ (start) $\mathbf{q}_1 = \mathbf{q}_t $
• $\mathbf{q}_2, \mathbf{q}_3$: Orientations at $t + \Delta t/2$ (midpoints): $\mathbf{q}_2 = \mathbf{q}_3 = \mathbf{q}_t \otimes \text{Exp}(\boldsymbol{\omega}_i \Delta t/2)$
• $\mathbf{q}_4$: Orientation at $t + \Delta t$ (end): $\mathbf{q}_4 = \mathbf{q}_t \otimes \text{Exp}(\boldsymbol{\omega}_i \Delta t)$

This allows the integrator to account for the change in the direction of the acceleration vector as the body rotates. By sampling the orientation mid-step, the filter accurately transforms body-frame acceleration into a curving world-frame trajectory.

Orientation Integration

Unlike position and velocity, which are additive, orientation integration uses the Closed-Form Rotation strategy within the RK4 framework.

• Slope Averaging: We average the angular velocities across the four stages. In the discrete case where $\boldsymbol{\omega}$ is constant, this simplifies to the unbiased measurement: $$\boldsymbol{\omega}_{avg} = \frac{1}{6}(\boldsymbol{\omega}_1 + 2\boldsymbol{\omega}_2 + 2\boldsymbol{\omega}_3 + \boldsymbol{\omega}_4) = \boldsymbol{\omega}_i$$
• Exponential Map: The orientation is updated using the right-multiplication convention (local frame):$$\mathbf{q}_{t+\Delta t} = \mathbf{q}_t \otimes \text{Exp}(\boldsymbol{\omega}_{avg} \Delta t) = \mathbf{q}_4$$ In implementation, this is performed via q * R.from_rotvec(omega * dt).
• Normalization: Following integration, the quaternion is re-normalized to counteract numerical drift and ensure it remains on the unit hypersphere:$$\mathbf{q} \leftarrow \frac{\mathbf{q}}{\|\mathbf{q}\|}$$

The Error State Jacobian ($\mathbf{F}_x$)

The transition matrix $\mathbf{F}_x$ (often denoted as $\Phi$ in classical navigation and in the Sola-appendix) defines how the error state $\delta \mathbf{x}$ evolves from time $t$ to $t + \Delta t$. While the nominal state is integrated using high-order RK4, the error state is propagated through a linearized system to maintain computational efficiency.

Continuous to Discrete Transformation

The error dynamics are governed by the continuous-time system $\delta \dot{\mathbf{x}} = \mathbf{A} \delta \mathbf{x}$ (compare Sola equ. 342). To implement this in software, we discretize the system using a first-order Taylor expansion of the matrix exponential (Solà Eq. 344):

$$\mathbf{F}_x = \exp(\mathbf{A} \Delta t) \approx \mathbf{I} + \mathbf{A} \Delta t$$

This linear truncation is the standard approximation. It assumes that the Jacobians are constant over the interval $\Delta t$.

The 18-State Discrete Transition Matrix

The Jacobian implemented in ESKFKinematics.get_error_jacobian corresponds to the block-wise truncation found in Solà Appendix C.2 (Eq. 401) resp.(Eq. 270). The state vector order is $[\delta\mathbf{p}, \delta\mathbf{v}, \delta\boldsymbol{\theta}, \delta\mathbf{a}_b, \delta\boldsymbol{\omega}_b, \delta\mathbf{g}]$.

$$\mathbf{F}_x = \begin{bmatrix} \mathbf{I} & \mathbf{I}\Delta t & \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0} \\ \mathbf{0} & \mathbf{I} & \mathbf{-R}\lfloor\mathbf{a}_u\rfloor_\times\Delta t & \mathbf{-R}\Delta t & \mathbf{0} & \mathbf{I}\Delta t \\ \mathbf{0} & \mathbf{0} & \mathbf{R}^T\{\boldsymbol{\omega}_u \Delta t\} & \mathbf{0} & -\mathbf{I}\Delta t & \mathbf{0} \\ \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{I} & \mathbf{0} & \mathbf{0} \\ \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{I} & \mathbf{0} \\ \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{I} \end{bmatrix}$$

Key Physical Interactions

• Velocity Row ($\mathbf{F}_{v\theta}$): The term $-\mathbf{R}\lfloor\mathbf{a}_u\rfloor_\times\Delta t$ couples orientation error to velocity. This is the most critical term for the tennis swing; it describes how a small tilt error incorrectly projects the heavy acceleration of the swing into the wrong world-frame axis.
• Orientation Row ($\mathbf{F}_{\theta\theta}$): The block $\mathbf{R}^T\{\boldsymbol{\omega}_u \Delta t\}$ represents the local-frame error propagation. It accounts for how the existing orientation uncertainty is "rotated" by the current angular velocity.
• Gravity ($\mathbf{F}_{vg}$): The $\mathbf{I}\Delta t$ block in the velocity row allows the filter to observe and correct the gravity vector estimate based on the integrated velocity residuals.

Numerical Hygiene

Because $\mathbf{F}_x$ is used to update the covariance matrix $\mathbf{P} \leftarrow \mathbf{F}_x \mathbf{P} \mathbf{F}_x^T + \mathbf{F}_i \mathbf{Q}_i \mathbf{F}_i^T$ (269), its structure must preserve the symmetry and positive-definiteness of $\mathbf{P}$. The blocks for biases ($\delta\mathbf{a}_b, \delta\boldsymbol{\omega}_b$) and gravity ($\delta\mathbf{g}$) are identity matrices, modeling these as random walks or constants.

Covariance of Perturbation Impulses

The covariance prediction is governed by the discrete-time transition matrix $\mathbf{F}_x$ and the noise Jacobian $\mathbf{F}_i$. Following the framework in Solà Section 5.4.3 (Equation 269), the covariance is updated as:

$$\mathbf{P} \leftarrow \mathbf{F}_x \mathbf{P} \mathbf{F}_x^T + \mathbf{F}_i \mathbf{Q}_i \mathbf{F}_i^T$$

with the assumption that the noise is isotropic (i.e. not rotation needed) see appendix E.2. of Sola.
$\mathbf{F}_i$ is simple a selection matrix (see Eq. 271 resp. Eq. 456).

Discrete Noise Covariance ($\mathbf{Q}_i$)

As specified in Solà Eq. 457, the diagonal matrix $\mathbf{Q}_i$ captures two distinct types of noise integration:

• Control Perturbations ($\sigma_{\tilde{\mathbf{a}}}, \sigma_{\tilde{\boldsymbol{\omega}}}$): Because the IMU measurements are treated as constant control signals over $\Delta t$ (ZOH), their associated errors are integrated as $\tilde{\mathbf{u}} \Delta t$. This leads to a variance scaling of $\Delta t^2$.
• Random Walks ($\sigma_{\mathbf{a}_w}, \sigma_{\boldsymbol{\omega}_w}$): The biases are modeled as Brownian motion. The uncertainty of a random walk grows linearly with time, leading to a variance scaling of $\Delta t$.

The Integrated Noise Matrix ($\mathbf{Q}_d$)

Assuming isotropic noise (where the rotation matrix $\mathbf{R}$ cancels out per Solà Appendix E.2), the 18x18 matrix $\mathbf{Q}_d = \mathbf{F}_i \mathbf{Q}_i \mathbf{F}_i^T$ is defined as:

$$\mathbf{Q}_d = \begin{bmatrix} \mathbf{0}_3 & \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0} \\ \mathbf{0} & \sigma_{\tilde{\mathbf{a}}}^2 \Delta t^2 \mathbf{I}_3 & \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0} \\ \mathbf{0} & \mathbf{0} & \sigma_{\tilde{\boldsymbol{\omega}}}^2 \Delta t^2 \mathbf{I}_3 & \mathbf{0} & \mathbf{0} & \mathbf{0} \\ \mathbf{0} & \mathbf{0} & \mathbf{0} & \sigma_{\mathbf{a}_w}^2 \Delta t \mathbf{I}_3 & \mathbf{0} & \mathbf{0} \\ \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0} & \sigma_{\boldsymbol{\omega}_w}^2 \Delta t \mathbf{I}_3 & \mathbf{0} \\ \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0} & \mathbf{0}_3 \end{bmatrix}$$

implemented in ESKFKinematics.get_process_noise

State Injection (apply_error_state)

The Error-State Kalman Filter maintains a nominal state ($\mathbf{x}$) and an error state ($\delta \mathbf{x}$). Once the update step calculates the optimal error $\delta \mathbf{x}$, it must be "injected" back into the nominal state.

Linear States

For position, velocity, and biases, injection is a simple addition: $$\mathbf{p} \leftarrow \mathbf{p} + \delta \mathbf{p}, \quad \mathbf{v} \leftarrow \mathbf{v} + \delta \mathbf{v}, \quad \mathbf{b} \leftarrow \mathbf{b} + \delta \mathbf{b}$$

Angular States (The Manifold Update)

We assume the error $\delta \boldsymbol{\theta}$ is defined in the local body frame:

$$\mathbf{q} \leftarrow \mathbf{q} \otimes \exp\left(\frac{\delta \boldsymbol{\theta}}{2}\right)$$

self.state.q * dq performs this right-multiplication, effectively rotating the current estimate by the error vector.

class ESKFKinematics:
    """Pure mathematical functions for ESKF physics."""

    @staticmethod
    def skew(v):
        """Returns the skew-symmetric matrix of a 3D vector."""
        return np.array([
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0]
        ])

    @staticmethod
    def get_derivatives(state: NominalState, imu: IMUInput):
        # Corrected measurements (Unbiased)
        a_u = imu.acc_raw - state.acc_bias 
        w_u = imu.gyro_raw - state.gyro_bias
        
        dp = state.v
        dv = state.q.apply(a_u) + state.g # R(q)*a_u + g
        return dp, dv, w_u

    
    @staticmethod
    def get_error_jacobian(state: NominalState, imu: IMUInput):
        """
        Computes the 18x18 discrete-time state transition matrix F_x.
        State order: [pos(0:3), vel(3:6), theta(6:9), ab(9:12), wb(12:15), g(15:18)]
        """
        dt = imu.dt
        F_x = np.eye(18)
        
        # 1. Pre-compute repeated terms
        R_matrix = state.q.as_matrix() # Rotation from Body to World
        a_u = imu.acc_raw - state.acc_bias
        w_u = imu.gyro_raw - state.gyro_bias
        
        # 2. Position Row
        # delta_p = delta_p + delta_v * dt
        F_x[0:3, 3:6] = np.eye(3) * dt
        
        # 3. Velocity Row
        # delta_v = delta_v + (-R*[a_u]_x * delta_theta - R*delta_ab + delta_g) * dt
        F_x[3:6, 6:9] = -R_matrix @ ESKFKinematics.skew(a_u) * dt
        F_x[3:6, 9:12] = -R_matrix * dt
        F_x[3:6, 15:18] = np.eye(3) * dt
        
        # 4. Orientation Row (Theta)
        # delta_theta = R{w_u*dt}^T * delta_theta - delta_wb * dt
        # The term R{w_u*dt} is the rotation increment
        rotation_increment = R.from_rotvec(w_u * dt).as_matrix()
        F_x[6:9, 6:9] = rotation_increment.T
        F_x[6:9, 12:15] = -np.eye(3) * dt
        
        # 5. Biases and Gravity Rows
        # (Remain Identity as they are modeled as random walks or constants)
        
        return F_x

    @staticmethod
    def get_process_noise(dt, sigma_acc, sigma_gyro, sigma_acc_bias, sigma_gyro_bias, sigma_g=0.0):
        """
        Computes the 18x18 process noise matrix Q_d.
        Strict implementation of Sola Eq. 456 & 457.
        """
        Q = np.zeros((18, 18))
    
        # 1. Velocity and Orientation noise (Control Perturbations)
        # Scaled by dt^2 per Sola Eq. 457 (Discrete-time perturbation model)
        v_var = (sigma_acc**2) * (dt**2)
        theta_var = (sigma_gyro**2) * (dt**2)
    
        Q[3:6, 3:6] = np.eye(3) * v_var
        Q[6:9, 6:9] = np.eye(3) * theta_var
    
        # 2. Bias Random Walk (Stochastic Perturbations)
        # Scaled by dt per Sola Eq. 457 (Wiener process/Random walk model)
        ab_var = (sigma_acc_bias**2) * dt
        wb_var = (sigma_gyro_bias**2) * dt
    
        Q[9:12, 9:12] = np.eye(3) * ab_var
        Q[12:15, 12:15] = np.eye(3) * wb_var
    
        # 3. Gravity Noise
        # Usually 0 in Sola, but can be set to a small value for numerical stability
        if sigma_g > 0:
            g_var = (sigma_g**2) * dt
            Q[15:18, 15:18] = np.eye(3) * g_var
        return Q

class IntegratorStrategy:
    def step(self, state: NominalState, imu: IMUInput, kinematics: ESKFKinematics):
        raise NotImplementedError

class EulerIntegrator(IntegratorStrategy):
    def step(self, state, imu, kinematics):
        dp, dv, w_u = kinematics.get_derivatives(state, imu)
        state.p += dp * imu.dt
        state.v += dv * imu.dt
        # Use a sub-strategy for rotation TODO
        return state

class RotationStrategy:
    def update(self, q: R, omega: np.ndarray, dt: float):
        raise NotImplementedError

class HigherOrderRotation(RotationStrategy):
    """
    Implements a higher-order integration of the rotation rates.
    
    """
    def update(self, q, omega, dt, omega_prev=None):
        if omega_prev is None:
            # Fallback to closed-form if no history
            return q * R.from_rotvec(omega * dt)
        
        # Calculate the angular acceleration (alpha) assuming linear change
        alpha = (omega - omega_prev) / dt
        
        # The 'Integrated' rotation vector over dt 
        # (This is a simplified 2nd order version of Solà's B.3)
        theta = (omega + 0.5 * alpha * dt) * dt
        
        return q * R.from_rotvec(theta)

Orientation Integration: Closed-Form Rotation

The orientation update is implemented in ClosedFormRotation using a multiplicative approach in the local frame. Unlike position and velocity, which use standard additive arithmetic, orientation belongs to the $SO(3)$ Lie Group, requiring the use of the Exponential Map.

class ClosedFormRotation(RotationStrategy):
    def update(self, q, omega, dt):
        return q * R.from_rotvec(omega * dt)

class DiscreteRK4Integrator(IntegratorStrategy):
    def step(self, state: NominalState, imu: IMUInput, kinematics: ESKFKinematics, rotator: RotationStrategy):
        dt = imu.dt
        
        # 1. Unbiased Measurements (Constant for this step)
        a_u = imu.acc_raw - state.acc_bias
        w_u = imu.gyro_raw - state.gyro_bias
        
        # 2. Intermediate Orientations (Used to rotate a_u)
        # These match q2_tmp, q3_tmp, and q4_tmp from your original code
        q1 = state.q
        q2 = rotator.update(q1, w_u, dt/2)
        q3 = rotator.update(q1, w_u, dt/2) # Identical to q2 in discrete case
        q4 = rotator.update(q1, w_u, dt)
        
        # 3. RK4 Slopes (k1, k2, k3, k4)
        # Each slope is a tuple of (dp, dv) -> (velocity, world_accel)
        
        # k1 (Start of interval)
        dp1 = state.v
        dv1 = q1.apply(a_u) + state.g
        
        # k2 (Midpoint using k1)
        dp2 = state.v + dv1 * (dt/2)
        dv2 = q2.apply(a_u) + state.g
        
        # k3 (Midpoint using k2)
        dp3 = state.v + dv2 * (dt/2)
        dv3 = q3.apply(a_u) + state.g
        
        # k4 (End of interval using k3)
        dp4 = state.v + dv3 * dt
        dv4 = q4.apply(a_u) + state.g
        
        # 4. Final Weighted Integration
        state.p += (dt / 6.0) * (dp1 + 2*dp2 + 2*dp3 + dp4)
        state.v += (dt / 6.0) * (dv1 + 2*dv2 + 2*dv3 + dv4)
        
        # Orientation update  
        state.q = q4
        # 5. Normalization
        q_norm = np.linalg.norm(state.q.as_quat())
        state.q = R.from_quat(state.q.as_quat() / q_norm)
        
        return state

class ESKF:
    def __init__(self, 
                 initial_state: NominalState, 
                 initial_covariance: np.ndarray, 
                 integrator: IntegratorStrategy,
                 rotator: RotationStrategy):
        self.state = initial_state
        self.P = initial_covariance
        
        # Dependency Injection
        self.integrator = integrator
        self.rotator = rotator
        self.kinematics = ESKFKinematics()

    def predict(self, imu: IMUInput, noise_params: dict):
        # Propagate Nominal State (The Physics)
        # Note: We pass the rotator strategy into the integrator step
        self.state = self.integrator.step(self.state, imu, self.kinematics, self.rotator)
        
        # Propagate Covariance (The Uncertainty)
        F_x = self.kinematics.get_error_jacobian(self.state, imu)
        Q = self.kinematics.get_process_noise(
            imu.dt, 
            noise_params['sigma_acc'], 
            noise_params['sigma_gyro'],
            noise_params['sigma_acc_bias'],
            noise_params['sigma_gyro_bias']
        )
        self.P = F_x @ self.P @ F_x.T + Q


    def apply_error_state(self, delta_x):
        """
        Injects the 18-element error state vector into the nominal state.
        delta_x: [dp, dv, dtheta, d_ab, d_wb, dg]
        """
        # 1. Position and Velocity (Standard Addition)
        self.state.p += delta_x[0:3]
        self.state.v += delta_x[3:6]
    
        # 2. Orientation (Quaternion Error Injection)
        # Solà Eq: q_true = q_est * Exp(d_theta)
        delta_theta = delta_x[6:9]
        if np.linalg.norm(delta_theta) > 1e-12:
            # Convert rotation vector to a quaternion
            # We use a right-multiplication for local error frames
            dq = R.from_rotvec(delta_theta)
            self.state.q = self.state.q * dq
    
        # 3. Biases (Standard Addition)
        self.state.acc_bias += delta_x[9:12]
        self.state.gyro_bias += delta_x[12:15]
    
        # 4. World Gravity (Standard Addition)
        # This matches the Jacobian H[:, 15:18] we defined earlier
        self.state.g += delta_x[15:18]


    def update_quiet_phase(self, imu_input, R_vel, R_g):
        """
        6-Element Update: Velocity (3) and Accel/Gravity (3).
        Ignores Gyro Bias to prevent 'Bias Leakage' from micro-motions.
        """
        # 1. Prepare Measurements
        y_acc = imu_input.acc_raw
        R_BW = self.state.q.inv().as_matrix()

        # 2. Predictions
        # h_acc: What the accelerometer SHOULD see if we are still (Gravity + Bias)
        h_acc = R_BW @ (-self.state.g) + self.state.acc_bias

        # 3. Innovations
        z_v = np.zeros(3) - self.state.v
        z_a = y_acc - h_acc
        z = np.concatenate([z_v, z_a]) # Length 6

        # 4. Construct 6x18 Jacobian H
        H = np.zeros((6, 18))
    
        # Velocity Innovation -> Velocity State (3:6)
        H[0:3, 3:6] = np.eye(3)

        # Accel Innovation -> Orientation (6:9), AccBias (9:12), Gravity (15:18)
        H[3:6, 6:9] = ESKFKinematics.skew(h_acc) # Tilt sensitivity
        H[3:6, 9:12] = np.eye(3)                # Accel bias sensitivity
        H[3:6, 15:18] = -R_BW                   # Gravity vector sensitivity

        # 5. Kalman Update
        R_block = block_diag(np.eye(3) * R_vel, np.eye(3) * R_g)
    
        S = H @ self.P @ H.T + R_block
        K = self.P @ H.T @ np.linalg.inv(S)

        delta_x = K @ z
        self.apply_error_state(delta_x)

        # Joseph Form for stability
        I = np.eye(18)
        IKH = I - K @ H
        self.P = IKH @ self.P @ IKH.T + K @ R_block @ K.T

    
    def update_stationary_with_g_estimation(self, imu_input, R_vel, R_g, R_gyro):
        """
        9-Element Update: Velocity (3), Accel/Gravity (3), Gyro/Bias (3)
        """
        # 1. Prepare Measurements
        y_acc = imu_input.acc_raw
        y_gyro = imu_input.gyro_raw
        R_BW = self.state.q.inv().as_matrix()
    
        # 2. Predictions (Nominal h(x))
        # We expect: 0 velocity, gravity offset by bias, and 0 rotation + bias
        h_acc = R_BW @ (-self.state.g) + self.state.acc_bias
        h_gyro = self.state.gyro_bias
    
        # 3. Innovations (z = measurement - prediction)
        z_v = np.zeros(3) - self.state.v
        z_a = y_acc - h_acc
        z_w = y_gyro - h_gyro # Any residual here is attributed to Gyro Bias
    
        z = np.concatenate([z_v, z_a, z_w]) # Now length 9
    
        # 4. Construct 9x18 Jacobian H
        H = np.zeros((9, 18))
        
        # Rows 0-2: Velocity Innovation -> Velocity State (3:6)
        H[0:3, 3:6] = np.eye(3)
    
        # Rows 3-5: Accel Innovation -> Tilt(6:9), AccBias(9:12), Gravity(15:18)
        H[3:6, 6:9] = ESKFKinematics.skew(h_acc)
        H[3:6, 9:12] = np.eye(3)
        H[3:6, 15:18] = -R_BW
    
        # Rows 6-8: Gyro Innovation -> Gyro Bias State (12:15)
        # Derivative of (y_gyro - gyro_bias) w.r.t gyro_bias is -Identity
        # But for the innovation z = y - h, where h = bias, d(z)/d(error_bias) is Identity
        H[6:9, 12:15] = np.eye(3)
    
        # 5. Kalman Gain and Covariance Update
        # Combine R values into a 9x9 block diagonal matrix
        R_block = block_diag(np.eye(3) * R_vel, np.eye(3) * R_g, np.eye(3) * R_gyro)
    
        S = H @ self.P @ H.T + R_block
        K = self.P @ H.T @ np.linalg.inv(S)
    
        delta_x = K @ z
        self.apply_error_state(delta_x)
    
        # Update P using Joseph Form
        I = np.eye(18)
        IKH = I - K @ H
        self.P = IKH @ self.P @ IKH.T + K @ R_block @ K.T

    

acc_tol=0.2
gyro_tol=(2*np.pi/360)*4

def detect_stationary(imu_input, g_mag=9.81, acc_tol=acc_tol, gyro_tol=gyro_tol):
    """
    Checks if the racket is likely stationary.
    acc_tol: m/s^2 deviation from g
    gyro_tol: rad/s threshold
    """
    acc_norm = np.linalg.norm(imu_input.acc_raw)
    gyro_norm = np.linalg.norm(imu_input.gyro_raw)
    is_still = (abs(acc_norm - g_mag) < acc_tol) and (gyro_norm < gyro_tol)
    #print (acc_norm, gyro_norm, is_still)
    return is_still

def run_predictions(imu_data, initial_state, 
                    noise_params, 
                    P_init,
                    r_vel, # velocity uncertainty if is nearly still
                    r_g, # g uncertainty if is nearly still
                    r_gyro, 
                    still_for): # pseudo-measurement if is_still for such time steps 
    """
    Test loop to evaluate dead-reckoning performance.
    imu_data: dict containing 'acc', 'gyro', and 't'
    initial_state: NominalState object
    noise_params: dict containing sigmas for acc, gyro, and biases
    """

    # We choose our strategies here (Dependency Injection)
    eskf = ESKF(
        initial_state=copy.deepcopy(initial_state),
        initial_covariance=P_init,
        #integrator=RK4Integrator(), 
        integrator=DiscreteRK4Integrator(),   
        rotator=ClosedFormRotation()     
    )
    
    # Storage for results
    results = []
    variances = [] # New list to store uncertainty

    # Stillness 
    R_vel = np.eye(3) * r_vel 
    R_g = np.eye(3) * r_g   # We trust the accelerometer for gravity, but not perfectly
    R_gyro = np.eye(3) * r_gyro
    still_count = 0
    # 2. The Main Loop (Simulating real-time sensor stream)
    for i in range(1, len(imu_data['t'])):
        dt = imu_data['t'][i] - imu_data['t'][i-1]
        
        # Create the Input Data Object
        imu_input = IMUInput(
            acc_raw=imu_data['acc'][i],
            gyro_raw=imu_data['gyro'][i],
            dt=dt
        )
        
        # The ESKF instance now handles everything internally
        # consistent with our modular "Orchestrator" pattern.
        eskf.predict(imu_input, noise_params)
        
        # Stillness Heuristic Detection
        is_still = detect_stationary(imu_input)
        if is_still:
            still_count += 1
            if still_count > still_for:
                # This pulls the 'g' state and 'v' state toward the heuristic
                eskf.update_stationary_with_g_estimation(imu_input, R_vel, R_g, R_gyro)
                #eskf.update_quiet_phase(imu_input, R_vel, R_g)
                #print("still", i)
        else:
            still_count = 0
        
        # Log the state for plotting
        results.append(eskf.state.to_array())
        # Extract the diagonal of P (the variances of the 18 error states)
        variances.append(np.diag(eskf.P).copy())

    
    return np.array(results), np.array(variances)

# 1. Define noise parameters (Adjust these to match your generate_imu_measurements call)
noise_params = {
    'sigma_acc': 0.01,         # Match imu_data acc_noise
    'sigma_gyro': 0.001,       # Match imu_data gyro_noise
    'sigma_acc_bias': 1e-4, # 0.0001, # Small value for prototype
    'sigma_gyro_bias': 1e-7 # 0.0001 # Small value for prototype
}

# Initialize the ESKF prediction uncertainty
P_init = np.eye(18) * 1e-8 # Baseline
# Position (0:3): Very sure. We define this as (0,0,0)
P_init[0:3, 0:3] = np.eye(3) * 1e-12
# Velocity (3:6): NOT sure. 
# Give it some room (e.g., .3 m/s uncertainty)
P_init[3:6, 3:6] = np.eye(3) * 0.03
# Tilt (6:9): Solà convention - Very sure (this is our reference frame)
P_init[6:9, 6:9] = np.eye(3) * 1e-12 # Or a very small epsilon like 1e-10
# Biases: Depends on your IMU datasheet
P_init[9:12, 9:12] = np.eye(3) * 0.1  # Accel Bias
P_init[12:15, 12:15] = np.eye(3) * 0.2 # Gyro Bias 
# High uncertainty in initial gravity vector
P_init[15:18, 15:18] = np.eye(3) * .1

r_vel = 5.e-1 # velocity uncertainty for stillness pseudo-measurement
r_g = .04 # gravity uncertainty for stillness pseudo-measurement
r_gyro = 1e-03
still_for = 5

# 1. Get the first IMU reading
first_acc = imu_data['acc'][0]

# 2. Rotate it to the World Frame using the initial orientation
# If acc = R.T * g, then g = R * acc
initial_q = R.from_quat(truth_data['quat'][0])
initial_g_estimate = initial_q.apply(-first_acc) # Note the negative sign

start_state = NominalState(
    p=truth_data['pos'][0], 
    v=truth_data['vel'][0], 
    q=initial_q,
    acc_bias=np.zeros(3), #np.zeros(3),
    gyro_bias=np.zeros(3),
    g=initial_g_estimate # gravity
)

# Run the filter 
results, variances = run_predictions(imu_data, start_state, noise_params, P_init, r_vel, r_g, r_gyro, still_for)
results.shape, variances.shape

((279, 19), (279, 18))

plot_eskf_results(truth_data, results, variances, imu_data, n_sigma=1)

plot_velocity_kinematics(truth_data, results, variances, imu_data, n_sigma=1)

plot_g_error(truth_data, results, variances, imu_data, n_sigma=1, true_g = G_GLOBAL)

plot_gyro_bias_error(results, variances, imu_data, true_bias=np.array([0.01, -0.02, 0.005]))