⚡️ Speed up method KalmanFilterXYAH.predict by 97%

[codeflash-ai]

📄 97% (0.97x) speedup for KalmanFilterXYAH.predict in ultralytics/trackers/utils/kalman_filter.py

⏱️ Runtime : 31.4 milliseconds → 15.9 milliseconds (best of 159 runs)

📝 Explanation and details

The optimized code achieves a 97% speedup by eliminating expensive NumPy operations and reducing memory allocations. The key optimizations are:

What was optimized:

1. Eliminated np.r_ concatenation: The original code used np.r_[std_pos, std_vel] which creates temporary arrays and performs concatenation. The optimized version pre-allocates a single np.empty(8) array and fills it directly with slice assignments.

2. Reduced redundant calculations: Instead of computing mean[3] multiple times (6 times in original), it's cached as h and reused in pos and vel calculations.

3. Replaced np.linalg.multi_dot: The original used the heavyweight multi_dot function for a simple 3-matrix chain. The optimized version breaks this into two separate @ operations, which is more efficient for this specific case.

4. Direct array operations: Replaced list creation and np.square() with direct element-wise multiplication (std_values * std_values).

Why it's faster:

• Line 44 in original (59.1% of runtime): np.diag(np.square(np.r_[std_pos, std_vel])) was the major bottleneck, involving list-to-array conversion, concatenation, squaring, and diagonal matrix creation. The optimized version reduces this to ~26% of total runtime.
• Memory efficiency: Fewer temporary arrays means better cache locality and reduced garbage collection overhead.
• Function call overhead: Eliminated expensive NumPy utility functions in favor of basic operations.

Impact on workloads:
The Kalman filter predict method is typically called in tight loops for multi-object tracking, where each frame processes dozens to hundreds of tracked objects. The 97% speedup directly translates to significant performance gains in real-time tracking applications, making the difference between meeting and missing frame rate requirements in computer vision pipelines.

Test case performance:
The optimization shows consistent 80-150% speedups across all test scenarios, from basic cases (single predictions) to large-scale stress tests (500-1000 tracks), indicating the optimization scales well with typical usage patterns.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 2960 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import time

import numpy as np

# imports
import pytest  # used for our unit tests
from ultralytics.trackers.utils.kalman_filter import KalmanFilterXYAH

# unit tests

# --------------------------- BASIC TEST CASES ---------------------------


def test_predict_identity_covariance():
    """Test predict with identity covariance and zero velocities."""
    kf = KalmanFilterXYAH()
    mean = np.array([10, 20, 1.5, 50, 0, 0, 0, 0], dtype=float)
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 32.8μs -> 18.1μs (80.9% faster)
    # Covariance should be positive semi-definite
    eigvals = np.linalg.eigvalsh(pred_cov)


def test_predict_nonzero_velocity():
    """Test predict with nonzero velocities."""
    kf = KalmanFilterXYAH()
    mean = np.array([10, 20, 1.5, 50, 1, -2, 0.5, 3], dtype=float)
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 30.3μs -> 15.0μs (103% faster)
    # Covariance should be positive semi-definite
    eigvals = np.linalg.eigvalsh(pred_cov)


def test_predict_different_height():
    """Test predict with different heights to check std_weight scaling."""
    kf = KalmanFilterXYAH()
    heights = [1, 10, 100, 1000]
    for h in heights:
        mean = np.array([0, 0, 1, h, 0, 0, 0, 0], dtype=float)
        covariance = np.eye(8)
        pred_mean, pred_cov = kf.predict(mean, covariance)  # 71.5μs -> 33.9μs (111% faster)
        # Motion covariance diagonal should scale with h in position and velocity
        expected_std_pos = kf._std_weight_position * h
        expected_std_vel = kf._std_weight_velocity * h


def test_predict_aspect_ratio_effect():
    """Test predict with varying aspect ratio, which should affect only mean[2]."""
    kf = KalmanFilterXYAH()
    mean = np.array([10, 20, 3.5, 50, 0, 0, 0, 0], dtype=float)
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 26.3μs -> 13.7μs (91.9% faster)


# --------------------------- EDGE TEST CASES ---------------------------


def test_predict_zero_height():
    """Test predict with zero height, should not crash and std_weight should be zero for position/velocity."""
    kf = KalmanFilterXYAH()
    mean = np.array([10, 20, 1.5, 0, 0, 0, 0, 0], dtype=float)
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 26.1μs -> 14.1μs (85.2% faster)


def test_predict_negative_height():
    """Test predict with negative height, should not crash and should scale std_weight accordingly."""
    kf = KalmanFilterXYAH()
    mean = np.array([10, 20, 1.5, -50, 0, 0, 0, 0], dtype=float)
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 26.6μs -> 13.5μs (96.2% faster)
    # std_weight_position and std_weight_velocity should be negative, but squared in motion_cov so positive
    expected_std_pos = kf._std_weight_position * mean[3]
    expected_std_vel = kf._std_weight_velocity * mean[3]


def test_predict_large_covariance():
    """Test predict with large covariance, should propagate uncertainty."""
    kf = KalmanFilterXYAH()
    mean = np.array([0, 0, 1, 10, 0, 0, 0, 0], dtype=float)
    covariance = np.eye(8) * 1e6
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 26.5μs -> 13.3μs (100% faster)


def test_predict_nan_inf_inputs():
    """Test predict with NaN and Inf values in mean and covariance."""
    kf = KalmanFilterXYAH()
    mean = np.array([np.nan, np.inf, 1, 10, 0, 0, 0, 0], dtype=float)
    covariance = np.eye(8)
    # Should propagate NaN/Inf to output mean
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 27.2μs -> 23.0μs (18.3% faster)


def test_predict_invalid_shapes():
    """Test predict with invalid input shapes, should raise an exception."""
    kf = KalmanFilterXYAH()
    mean = np.array([1, 2, 3, 4, 5, 6, 7])  # Only 7 elements
    covariance = np.eye(8)
    with pytest.raises(ValueError):
        kf.predict(mean, covariance)  # 34.2μs -> 17.8μs (92.2% faster)
    mean = np.zeros(8)
    covariance = np.eye(7)
    with pytest.raises(ValueError):
        kf.predict(mean, covariance)  # 20.7μs -> 9.21μs (125% faster)


def test_predict_non_float_inputs():
    """Test predict with integer inputs, should work and output float arrays."""
    kf = KalmanFilterXYAH()
    mean = np.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=int)
    covariance = np.eye(8, dtype=int)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 38.9μs -> 23.8μs (62.9% faster)


# --------------------------- LARGE SCALE TEST CASES ---------------------------


def test_predict_many_tracks():
    """Test predict on a large batch of tracks."""
    kf = KalmanFilterXYAH()
    n = 500  # Large but <1000
    means = np.tile(np.array([10, 20, 1.5, 50, 1, -2, 0.5, 3], dtype=float), (n, 1))
    covariances = np.tile(np.eye(8), (n, 1, 1))
    results_mean = []
    results_cov = []
    for i in range(n):
        pred_mean, pred_cov = kf.predict(means[i], covariances[i])  # 5.30ms -> 2.76ms (92.2% faster)
        results_mean.append(pred_mean)
        results_cov.append(pred_cov)
    # The first result should match the single-track test
    single_pred_mean, single_pred_cov = kf.predict(means[0], covariances[0])  # 11.3μs -> 5.59μs (102% faster)


def test_predict_performance_large_batch():
    """Test performance for a large batch (timing should be reasonable)."""
    kf = KalmanFilterXYAH()
    n = 800
    means = np.tile(np.array([10, 20, 1.5, 50, 1, -2, 0.5, 3], dtype=float), (n, 1))
    covariances = np.tile(np.eye(8), (n, 1, 1))
    start = time.time()
    for i in range(n):
        kf.predict(means[i], covariances[i])  # 8.32ms -> 4.23ms (96.8% faster)
    elapsed = time.time() - start


def test_predict_stress_randomized():
    """Test predict under randomized inputs for robustness."""
    kf = KalmanFilterXYAH()
    rng = np.random.default_rng(42)
    for _ in range(100):
        mean = rng.normal(loc=0, scale=100, size=8)
        covariance = rng.normal(loc=0, scale=10, size=(8, 8))
        covariance = covariance @ covariance.T + np.eye(8) * 1e-6  # Make positive definite
        pred_mean, pred_cov = kf.predict(mean, covariance)  # 1.35ms -> 547μs (146% faster)
        eigvals = np.linalg.eigvalsh(pred_cov)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import numpy as np

# imports
import pytest  # used for our unit tests
from ultralytics.trackers.utils.kalman_filter import KalmanFilterXYAH

# unit tests

# ----------- BASIC TEST CASES -----------


def test_predict_identity_covariance():
    # Basic test: mean at origin, identity covariance
    kf = KalmanFilterXYAH()
    mean = np.array([0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0])
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 30.6μs -> 15.7μs (94.6% faster)


def test_predict_nonzero_velocity():
    # Basic test: mean with nonzero velocity
    kf = KalmanFilterXYAH()
    mean = np.array([10.0, 20.0, 2.0, 5.0, 1.0, -1.0, 0.5, 2.0])
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 27.4μs -> 14.3μs (91.8% faster)
    # The first four elements should be updated by velocity
    expected_mean = mean.copy()
    expected_mean[:4] += mean[4:]


def test_predict_with_non_identity_covariance():
    # Basic test: non-identity covariance
    kf = KalmanFilterXYAH()
    mean = np.array([5.0, 5.0, 1.0, 10.0, 0.0, 0.0, 0.0, 0.0])
    covariance = np.eye(8) * 2
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 26.3μs -> 13.1μs (100% faster)


def test_predict_with_negative_velocity():
    # Basic test: negative velocity
    kf = KalmanFilterXYAH()
    mean = np.array([10.0, 10.0, 1.0, 10.0, -2.0, -2.0, -0.5, -1.0])
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 27.4μs -> 15.4μs (78.4% faster)
    # Position should decrease by velocity amount
    expected_mean = mean.copy()
    expected_mean[:4] += mean[4:]


# ----------- EDGE TEST CASES -----------


def test_predict_zero_height():
    # Edge case: zero height
    kf = KalmanFilterXYAH()
    mean = np.array([0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 30.5μs -> 16.1μs (89.9% faster)


def test_predict_large_height():
    # Edge case: very large height
    kf = KalmanFilterXYAH()
    large_h = 1e6
    mean = np.array([0.0, 0.0, 1.0, large_h, 0.0, 0.0, 0.0, 0.0])
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 28.9μs -> 15.6μs (86.0% faster)


def test_predict_negative_height():
    # Edge case: negative height (physically invalid, but should not crash)
    kf = KalmanFilterXYAH()
    mean = np.array([0.0, 0.0, 1.0, -10.0, 0.0, 0.0, 0.0, 0.0])
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 29.0μs -> 15.4μs (87.7% faster)


def test_predict_nan_input():
    # Edge case: input contains NaN
    kf = KalmanFilterXYAH()
    mean = np.array([np.nan, 0.0, 1.0, 10.0, 0.0, 0.0, 0.0, 0.0])
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 28.3μs -> 15.1μs (87.0% faster)


def test_predict_inf_input():
    # Edge case: input contains inf
    kf = KalmanFilterXYAH()
    mean = np.array([np.inf, 0.0, 1.0, 10.0, 0.0, 0.0, 0.0, 0.0])
    covariance = np.eye(8)
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 29.5μs -> 23.1μs (27.4% faster)


def test_predict_singular_covariance():
    # Edge case: singular covariance matrix (all zeros)
    kf = KalmanFilterXYAH()
    mean = np.array([0.0, 0.0, 1.0, 10.0, 0.0, 0.0, 0.0, 0.0])
    covariance = np.zeros((8, 8))
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 28.7μs -> 15.9μs (80.5% faster)
    # Covariance should be equal to motion_cov
    std_pos = [
        kf._std_weight_position * mean[3],
        kf._std_weight_position * mean[3],
        1e-2,
        kf._std_weight_position * mean[3],
    ]
    std_vel = [
        kf._std_weight_velocity * mean[3],
        kf._std_weight_velocity * mean[3],
        1e-5,
        kf._std_weight_velocity * mean[3],
    ]
    motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))


def test_predict_non_square_covariance():
    # Edge case: non-square covariance should raise an error
    kf = KalmanFilterXYAH()
    mean = np.zeros(8)
    covariance = np.zeros((8, 7))  # Not square
    with pytest.raises(ValueError):
        kf.predict(mean, covariance)  # 29.7μs -> 16.3μs (82.6% faster)


def test_predict_wrong_shape_mean():
    # Edge case: mean vector wrong shape should raise an error
    kf = KalmanFilterXYAH()
    mean = np.zeros(7)
    covariance = np.eye(8)
    with pytest.raises(IndexError):
        kf.predict(mean, covariance)


def test_predict_wrong_shape_covariance():
    # Edge case: covariance matrix wrong shape should raise an error
    kf = KalmanFilterXYAH()
    mean = np.zeros(8)
    covariance = np.eye(7)
    with pytest.raises(ValueError):
        kf.predict(mean, covariance)  # 44.4μs -> 25.2μs (76.4% faster)


# ----------- LARGE SCALE TEST CASES -----------


def test_predict_large_number_of_tracks():
    # Large scale: predict for many tracks in a loop
    kf = KalmanFilterXYAH()
    n_tracks = 500  # keep under 1000 as per instructions
    means = np.random.randn(n_tracks, 8) * 10
    covariances = np.array([np.eye(8) for _ in range(n_tracks)])
    for i in range(n_tracks):
        pred_mean, pred_cov = kf.predict(means[i], covariances[i])  # 5.25ms -> 2.67ms (96.5% faster)


def test_predict_performance_large_batch():
    # Large scale: predict for a large batch and check time (not strict, just that it runs)
    kf = KalmanFilterXYAH()
    n_tracks = 999
    means = np.random.randn(n_tracks, 8) * 100
    covariances = np.array([np.eye(8) * 5 for _ in range(n_tracks)])
    # Just run and check output shape
    for i in range(n_tracks):
        pred_mean, pred_cov = kf.predict(means[i], covariances[i])  # 10.4ms -> 5.29ms (96.7% faster)


def test_predict_extreme_values():
    # Large scale: extreme values in mean and covariance
    kf = KalmanFilterXYAH()
    mean = np.array([1e9, -1e9, 1e5, 1e8, 1e7, -1e7, 1e6, -1e6])
    covariance = np.eye(8) * 1e12
    pred_mean, pred_cov = kf.predict(mean, covariance)  # 32.9μs -> 18.1μs (82.3% faster)


def test_predict_multiple_calls_consistency():
    # Large scale: multiple calls with the same input should yield the same output
    kf = KalmanFilterXYAH()
    mean = np.array([5.0, 5.0, 2.0, 20.0, 1.0, 1.0, 0.5, 2.0])
    covariance = np.eye(8) * 3
    pred_mean1, pred_cov1 = kf.predict(mean, covariance)  # 28.0μs -> 13.9μs (101% faster)
    pred_mean2, pred_cov2 = kf.predict(mean, covariance)  # 16.2μs -> 7.04μs (130% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-KalmanFilterXYAH.predict-mir8nwhp and push.