Add sparse hessian to deal with multiple or incomplete observations

[zhengp0]

Allow number of rows in data to be different than number of variables

• create design matrix using jax.experimental.sparse.BCOO
• replace the original hessian_diag function by hessian in the likelihood (this will influence model.hessian and preconditioner)
• added a extra argument called data_span in the function KernelRegModel.fit, in case data is incomplete and does not contain all coordinate we are interested.
• added a predict function for convenient prediction

Note

• Current BCOO seems cannot handle matrix-matrix multiplication, at least not on my computer. We need to create a operator for Hessian in the likelihood, which is probably preferable anyway :)
• We cannot jit the function that will return a function, please see New TypeError: <class 'function'> is not a valid JAX type exception in JAX 0.2.0 jax-ml/jax#4416, and all the objective, gradient and hessian computation are vectorized, so maybe we don't need jit for them.

[AHsu98]

We should meet about this. I think we can generically handle multiple observations by wrapping into a modified likelihood. The only time we have a non diagonal hessian for the likelihood is when it involves a sum of multiple predictions (integral or something). I’m also about to push some minor changes I ran into from doing the sub national estimation from Reed.

[zhengp0]

We should meet about this. I think we can generically handle multiple observations by wrapping into a modified likelihood. The only time we have a non diagonal hessian for the likelihood is when it involves a sum of multiple predictions (integral or something). I’m also about to push some minor changes I ran into from doing the sub national estimation from Reed.

That sounds good! I will send an invite, sorry just saw this. I will leave this PR open until we discuss.