Error estimate in "Hierarchical Bayesian Linear Regression" example

[leostenzel]

Hi, I'm still learning how to use this cool package:
Working through the Hierarchical Bayesian Linear Regression example, I was wondering about the computation of the prediction's error.

Comparing with the numpyro notebook, the error of FVC seems a bit off; the values are reasonable, but the smallest variance is always at week 0.
I understand this is coming from computing the variance as var(patient_α) + var(patient_β) * week ^ 2, and I think I get the formula, but that doesn't always seem to match the data?
In particular, if you offset the "Weeks" in the input, the predicted error changes, and that does not seem to be the case with numpyro.

Is there any way you can fix this in RxInfer, or should you transform your input, or…?
It's not an urgent question, just trying to understand RxInfer's capabilities and limitations.

Edit: I guess I could try to be clearer. Taking a patient who isn't already shown in the examples, ID00042637202184406822975, the original input data give in RxInfer (left) and numpyro (right) the prediction:

Both seem reasonable, but intuitively I'd think that the minimal variance should be where the most measurements are, i.e., around week 20, as predicted by numpyro.

Things get worse when I shift the data by, e.g., -100 weeks.
I have to admit that the numpyro result changes more than I first realized. I guess that could just be because I did not adjust the priors?
But RxInfer's prediction (left) is clearly worse, and variance just doesn't make any sense.

I suppose this got something to do with patient_α and patient_β not being independent?
Maybe that doesn't matter too much for the overall prediction…
Just, reading the tutorial, I don't know if I agree with "The results match our expectations perfectly!".
And I think it could be quite instructive if you could elaborate on the assumptions, and how one can estimate the error more accurately.
Thanks, Leo

[albertpod]

Hi @leostenzel! Thanks for the question! We are a little busy at the moment, but we noted your problem.
We will try to reply in the thread within a week ;)

[albertpod]

@leostenzel I've reviewed this nb once more. It seems to me that var(patient_α) + var(patient_β) * week ^ 2 is incorrect. The predictive distribution should take into account the variance (inverse of σ) coming from the likelihood, i.e.

likelihood_var = inv(mean(results.posteriors[:σ]))
# rest of the code
pv = var(patient_α) + var(patient_β) * week ^ 2 + likelihood_var

Resulting in higher variance in predictions:

@bartvanerp @ismailsenoz , would you agree with that?

I also think whether we can include the missing observations inside the model so we can generate predictions during the inference without having this extra piece of code as shown here. Writing down the formula for predictions manually is not user friendly.
Perhaps @bvdmitri can check that once he's back.

@leostenzel, as for the shift, can you explain a bit more what you mean that you've shifted data by -100 weeks?

[bartvanerp]

The difference in error estimates is a consequence of the methods that we employ, which cannot be compared directly to Numpyro. Numpyro is fully sampling-based and can converge to the exact posterior estimates if left running for a long enough time. For the current problem where there are no time constraints, Numpyro will be a suitable tool. However, for other problems that rely on real-time signal processing, sampling-based methods simply cannot keep up and this is exactly the area where RxInfer shines. You are right though that we are compromising accuracy here, which is something that we are still working on to improve.

[albertpod]

@leostenzel, if you'd like to include bias into RxInfer.jl model, let me know.
I am curious about the speed of numpyro for this example; how does it compare to RxInfer.jl?

[leostenzel]

@bartvanerp Absolutely, RxInfer is much faster, and seems really cool. It's just about understanding what does and doesn't work.

@albertpod On a single core of my Apple M1, I see about 2min for numpyro (4000+4000 samples seem converged; a bit longer with t_offset), and about 2-5s with RxInfer (for 50 to 100 iterations).
Could be interesting to improve the RxInfer model… Just seems more important to get the error estimate right. Maybe using missing data to compute predictions already does that?

[albertpod]

Thanks, @leostenzel, for the update.

Let me sketch how to proceed with allow_missing.

First, in the model specification, we change:
data = datavar(Int, n_codes)
to
data = datavar(Int, n_codes) where {allow_missing = true} .
We must then change our dataset to accommodate the weeks with missing FVC. For example,

function add_missing(dataset, week_range)
    new_dataset = DataFrame(Patient = String[], Weeks = Int[], FVC = Union{Int, Missing}[])

    for id in unique(dataset.Patient)
        id_weeks = filter(row -> row.Patient == id, dataset).Weeks
        existing_fvc = filter(row -> row.Patient == id, dataset).FVC

        for week in week_range
            fvc = week in id_weeks ? existing_fvc[findfirst(x -> x == week, id_weeks)] : missing
            push!(new_dataset, (id, week, fvc))
        end
    end

    return new_dataset
end

Then

weeks_range = -12:130

new_dataset = add_missing(dataset, weeks_range)

patient_ids          = new_dataset[!, "Patient"] # get the column of all patients
patient_code_encoder = Dict(map(((id, patient), ) -> patient => id, enumerate(unique(patient_ids))));
patient_code_column  = map(patient -> patient_code_encoder[patient], patient_ids)

new_dataset[!, :PatientCode] = patient_code_column

first(patient_code_encoder, 5)

As we have allowed missing the observations, RxInfer.jl will treat them as random variables hence we need to include the data in constraints:

@constraints function partially_pooled_constraints()
    # rest of the code
    q(FVC_est, data, σ) = q(FVC_est, data)q(σ) # you need structured between FVC_est and data
end 

Then, in theory, you can simply run and extract the predictions for every patient

partially_pooled_inference_results = partially_pooled_inference(new_dataset)

id = "ID00009637202177434476278"
selected_rows = findall(row -> row.Patient == id, eachrow(new_dataset))
predicted_mean = mean.(partially_pooled_inference_results.predictions[:data][selected_rows])
predicted_std = std.(partially_pooled_inference_results.predictions[:data][selected_rows])

p = patientchart(dataset, id)
plot!(weeks_range, predicted_mean, ribbon = predicted_std, color = :orange, label = "Predicted FVC")

The bad news is that there's apparently a bug with predictions functionality, so this won't work with the public version of RxInfer.jl. I need to double-check and open the issue, @bvdmitri.

The error I am getting when using structured constraints on the likelihood is:
MethodError: no method matching [/](https://file+.vscode-resource.vscode-cdn.net/)(::PointMass{Int64}, ::Int64)

To run in with the public version of RxInfer.jl you need to change the constraints to mean-field:

@constraints function partially_pooled_constraints()
    # rest of the code
    q(FVC_est, data, σ) = q(FVC_est)q(data)q(σ) # This is bad as we won't take into account the variance from FVC_est RV.
end 

So the last option won't provide good variance estimate. Anyway, I will double check the issue with a simpler model and submit the issue if needed tomorrow. Thanks for trying out!

EDIT:
actually

q(FVC_est)q(data)q(σ)

is a correct specification in this case, but the predictive variance is estimated incorrectly. I am investigating it.

[leostenzel]

Wow, thanks for investigating and writing all this up 🙏
It's cool to have RxInfer.jl compute the prediction and variance right away; avoiding errors when I try myself…

But now I get partially_pooled_inference_results.predictions[:data] .|> std |> allequal # true, that's not very helpful?
I might have made a mistake somewhere, but can't seem to find it.

It's cool that I can now directly plot mean and std of partially_pooled_inference_results.posteriors[:FVC_est].
But it's exactly the same as variance as var(patient_α) + var(patient_β) * week ^ 2 gave me…

Don't stress about it over the holidays! It's not at all urgent for me; I'm just curious if there's a more accurate estimate.

[albertpod]

Hi @leostenzel. I'm sorry for not getting back to you sooner.
So, partially_pooled_inference_results.predictions[:data] .|> std |> allequal # true is associated with this issue now.

As for var(patient_α) + var(patient_β) * week ^ 2, this is indeed incorrect and must be fixed in the notebook; we have to add inv(mean(results.posteriors[:σ])) to the variance estimate. I believe the resulting variance will differ from numpyro, but this is because of the differences in inference methods.

In principle, this should give the same result as I was hoping to get when performing predictions within the inference function, but a "bug" has been detected.