ENH: add spherical model kernel

[esavary]

• Integrate custom kernel based on spherical covariance function.
• Add tests for the kernel.

[oesteban]

I had a very quick pass and left a few suggestions.

[oesteban]

This doesn't document the argument, indeed, it is the second argument to this function. Please describe what is the effect of this argument and indicate whether it is not used.

[esavary]

Let's first discuss with @jhlegarreta about moving the angle computation inside. From this, we will know precisely what Y is

[jhlegarreta]

about moving the angle computation inside

Talked to Óscar yesterday about this, and we concluded that my initial thought of providing the kernel with the angles is probably not right. So I take the blame for having pushed things into that direction.

Óscar mentioned that this would allow the GP to do the below computation appropriately when providing the query bvec in the prediction:
https://github.com/scikit-learn/scikit-learn/blob/972e17fe1aa12d481b120ad4a3dc076bae736931/sklearn/gaussian_process/_gpr.py#L439

So the kernel should receive the bvecs, and make the computation of the angles just before these lines:
https://github.com/nipreps/eddymotion/pull/203/files#diff-7401613f4c745123dea57ac1f74163b4e14f6ed2e0c82fa2e466835139743fdbR99-R100 and https://github.com/nipreps/eddymotion/pull/203/files#diff-7401613f4c745123dea57ac1f74163b4e14f6ed2e0c82fa2e466835139743fdbR103-R105

e.g.

X = compute_pairwise_angles(bvecs.T, closest_polarity=True)
K = self.lambda_s * np.where(
    X <= self.a, 1 - 3 * (X / self.a) ** 2 + 2 * (X / self.a) ** 3, 0
) + self.sigma_sq * np.eye(len(X))

Doing the above in an example with a simulated signal (using https://github.com/nipreps/eddymotion/blob/a8742d955d9a43676d54ef194737a42db6f74015/docs/notebooks/dwi_simulated_gp.ipynb) triggers a dimensionality error on K_gradient by the underlying call
https://github.com/scikit-learn/scikit-learn/blob/972e17fe1aa12d481b120ad4a3dc076bae736931/sklearn/gaussian_process/_gpr.py#L297

So the computation of K_gradient will need to be adapted now.

The dimensionality issues that you talked about in other messages is probably related to this. Fixing the above will hopefully will put us in the right track.

[jhlegarreta]

Thanks for doing this @esavary.

Once the two outstanding comments are solved, this is ready to go IMO.

Further improvements (including my own comments) can be addressed in separate PRs.

[jhlegarreta]

So I see that in https://scikit-learn.org/stable/modules/generated/sklearn.gaussian_process.kernels.Kernel.html they do the X=Y assignment if Y is None (note the absence of the not).

What if we avoid doing the X=Y assignment to get past this situation? My thought is linked to
This is related to https://github.com/nipreps/eddymotion/pull/203/files#r1641440923, and I now completely see your point in #200 (comment) and #200 (comment) @esavary:

if Y is None:
  theta = compute_pairwise_angles(X, X)
else:
  theta = compute_pairwise_angles(X, Y)

 K = self.lambda_s * np.where(
                theta <= self.a, 1 - 3 * (theta / self.a) ** 2 + 2 * (theta / self.a) ** 3, 0
            )

(assuming here X and potentially Y are a set of gradient vectors provided by the caller.)

(Have not checked whether we would still need to multiply by np.eye(len(theta) if Y is None)

The eval_gradient block may need to be adjusted concerning the dimensions now that X is assumed to be bvecs.

So compute_pairwise_angles would need to be modified to accept two input parameters. I apologize for not having realized about the downstream needs and implications of these two aspects. Really sorry @esavary.

[esavary]

So compute_pairwise_angles would need to be modified to accept two input parameters. I apologize for not having realized about the downstream needs and implications of these two aspects. Really sorry @esavary.

No worries, it also makes a lot of sense to compute the angles beforehand since we have all the information at the beginning. It would just have required us to recode the fit method. I'll suggest a generalization of the angle computation in a separate PR and then use this to update the kernel to take the bvec directly. This will also solve the problem of X=Y because we can then perform different computations when Y is given.

[jhlegarreta]

PR #212 to make compute_pairwise_angles accept two parameters.

[esavary]

I turned the PR into a draft for the moment because it depends on PR #212 . I integrated the angle computation into the kernel using the new function from that PR.

[jhlegarreta]

We probably want to call this
https://github.com/nipreps/eddymotion/pull/213/files#diff-ba53cdfdd502439a31a24c148f9be58dab5fb56d09b25547557933b3b259e46fR57

Making that a function will allow for reusing it in the multi-shell case. Also, will allow us to ensure that the behavior is the expected one prior to the kernel computations (e.g. together with PR #214).

[jhlegarreta]

This can probably be simplified since the pairwise angle computation now accepts the second argument to be None. Also, looking at $\mathbf{K}_{y} = \mathbf{K} + \sigma^2 \mathbf{I}$ just below eq. 12 in Andresson15 maybe self.sigma_sq * np.eye(len(angles)) needs to be added in either case, and thus we would not need the if/else statement.