Add factorial API

[SamDuffield]

My first attempt at a factorial API. I appreciate it's quite a lot of code added (and this is just to support Gaussian methods). Happy to break it up into multiple PRs if we think that's easier to grasp and there's a logical way to do so.

In terms of modifications to the existing cuthbert code I think they are minimal. I just had to change the init_prepare in Kalman so that it doesn't actually touch the observation parameters if the observation y is all nans. This is to avoid mismatching dimensions because in this factorial API the initial step is very different (acts globally across factors) to the subsequent steps (acts locally only on selected factors). (I think we'll also have to change this for the other methods too)

If anything is unclear please shout and I can document better!

[Sahel13]

Do we want to support GetFactorialIndices returning an integer instead of an array? Some of this logic breaks if we pass a scalar or a 0-dimensional array.

[SamDuffield]

Do we want to support GetFactorialIndices returning an integer instead of an array? Some of this logic breaks if we pass a scalar or a 0-dimensional array.

Yes I think we do as this is useful for smoothing and potentially other things later (like predict). It's a bit tricky as I think for filtering we want it to be strictly 1-dimensional but more generally it can be 0-dim

[Sahel13]

Yes I think we do as this is useful for smoothing and potentially other things later (like predict). It's a bit tricky as I think for filtering we want it to be strictly 1-dimensional but more generally it can be 0-dim

Then I think the easiest fix is to always have a factorial dimension, even if factorial_inds is a single integer. At the moment _join_arr would treat a 2-dimensional covariance matrix as a vector of means, and would flatten the whole thing.

[SamDuffield]

So factorial_inds is only used in Extract and Insert. We should add some code that makes sure it handles both cases ndim=0 and ndim=1 robustly.

Think about whether e.g. array([0]) should keep the factorial dimension or not (array(0) should definitely remove it). Probably should keep so there is an option to keep that dimension with length one.

[SamDuffield]

@Sahel13 I've updated the docstrings with the above comments. I think factorial.gaussian.extract and factorial.gaussian.insert already had the desired behaviour. The issue you mentioned regarding _join_arr would only occur if the input to join did not have a factorial dimension, which shouldn't happen. We could add some checks for this in join but it might be a bit messy

[Sahel13]

@Sahel13 I've updated the docstrings with the above comments. I think factorial.gaussian.extract and factorial.gaussian.insert already had the desired behaviour. The issue you mentioned regarding _join_arr would only occur if the input to join did not have a factorial dimension, which shouldn't happen. We could add some checks for this in join but it might be a bit messy

Is there something wrong with this test then:

def test_extract_and_join_preserves_single_factor_covariance_for_scalar_index():
    model_params = generate_factorial_kalman_model(
        seed=1,
        x_dim=2,
        y_dim=1,
        num_factors=3,
        num_factors_local=1,
        num_time_steps=1,
    )
    filter_obj, _, model_inputs = build_pairwise_factorial_filter(model_params)
    factorializer = factorial.gaussian.build_factorializer(
        get_factorial_indices=lambda _: 1
    )

    init_state = filter_obj.init_prepare(model_inputs[0])
    joined_state = factorializer.extract_and_join(init_state, model_inputs[1])

    assert joined_state.chol_cov.shape == (2, 2)

[SamDuffield]

Is there something wrong with this test then:

Kinda, your factorial_inds is an int (i.e. ndim=0) but is being sent to join (via extract_and_join). But join requires a factorial dimension. I've added an assert statement to extract_and_join to check this

[Sahel13]

I'm a bit confused now, so we don't actually support GetFactorialIndices returning a scalar? Because cuthbert.factorial.filtering.filter doesn't work with that now since there is an extract_and_join inside the body_local function.

If this is intended behavior, I'm happy to merge this.

[SamDuffield]

Scalar factorial_inds is not supported for filtering but is supported when using extract on its own which is useful for smoothing and perhaps predict (which is not supported yet)

I agree it's a bit confusing, not sure how we can document better