Fix ordering of time-varying statespace matices for UnivariateFilter's kalman step

[sebcroft]

This aims to address #612

Following the error in #612, the design matrix Z should have shape (2,3) not (3,3). Intial thought is the latter shape corresponds to the transition matrix T.

In BaseFilter, the kalman_step function calls the method unpack_args to reorganise the inputs into a standardized order. This is needed because when the user specifies time varying matrices, the kalman_step arguments (y, a, P, c, d, T, Z, R, H, Q) can get 'mixed-up' depending on which variables are assigned to the sequences and non_sequences in pytensor.scan.

UnivariateFilter redefines the kalman_step function and does not currently call unpack_args. Therefore, when the selection matrix R is the only time varying matrix (as in #612) it gets sent to the sequences/start (with y) and shifts the remaining args up by 1. As a result, T is incorrectly used for the Z input, which aligns with the shape error.

When all matrices are time-invariant unpack_args is not required because the correct ordering is preserved -- this is why UnivariateFilter works for the time-invariant matrices as mentioned in #612.

To address this I have added the method unpack_args to the kalman_step in UnivariateFilter.

Some extra stuff:

The above changes should hopefully be enough, but to keep things consistent with the BaseFilter you could possibly deal with the nan values using handle_missing_values instead i.e.:

def kalman_step(self, *args):
    
    y, a, P, c, d, T, Z, R, H, Q = self.unpack_args(args)
    
    nan_mask = pt.isnan(y)
    
    y_masked, Z_masked, H_masked, all_nan_flag = self.handle_missing_values(y, Z, H)
    
    result = pytensor.scan(
        self._univariate_inner_filter_step,
        sequences=[y_masked, Z_masked, d, pt.diag(H_masked), nan_mask],
        outputs_info=[a, P, None, None, None],
        name="univariate_inner_scan",
        return_updates=False,
    )

I'm happy to add this if you think it's best!

[jessegrabowski]

Great catch -- thanks for this!

I think it makes sense to also add the missing value handling if you don't object.

Also make sure you install and run pre-commit in your local dev environment. It's just pip install pre-commit; pre-commit install; pre-commit run --all (after you do install it will automatically run each time you git commit)