There’s a way you can instrument the code to see where the compute’s going:
Unfolding Kronecker products is rough. If there’s any way you can keep this implicit, it’ll probably be an enormous help. Especially since you’re just Kroneckering with the identity matrix in function A_sharp.
You don’t need parens around the return.
You can declare and define at the same time so this can be
int m = rows(A);
This is a very simplified product. I would just go through and define the entries of A_temp directly here and not try to do them in terms of big identity matrices.
There is almost never a justification for creating an identity matrix. You use this as matrix_exp(A_sharp_temp * ts) - I. This is better coded in Stan as
add_diag(matrix_exp(A_sharp_temp * ts), -1);
It will broadcast the -1 for you so you don’t need to allocate a vector.
This is just going to be nasty. Given it has a Kronecker structure, I would work with the Kronecker inverse, (A \otimes B)^{-1} = A^{-1} \otimes B^{-1}. But then the Kronecker structure from A_sharp isn’t exactly this. But whatever it is, simplifying this should be a top priority.
Given the complexity of this model, I’d urge you to put constraints where relevant on the data just for a sanity check and to help readers understand it.
Anything that’s constant should be declared in a transformed data block—this will be a big savings on useless autodiff that’d otherwise happen.
For filling the upper triangle, you can initialize to 0:
matrix[J,J] A = rep_matrix(0, J, J);
Then to fill, you want to create a 2d matrix
transformed data {
array[2, num_effects - J] idxs = ...; // this should be pairs of (i, j) where effects_mat[i, j] == 1
}
...
for (j in 1:num_effects - J) {
A[idxs[j, 1], idxs[j, 2]] = A_offdiag[j];
}
It reduces a quadratic loop to sublinear.
this can just be z_drift[i-1]', but it’s even better if you represent z_drift as an array of vectors rather than as a matrix—it saves an extra allocation and copy.
You don’t ever want to invert and multiply. It’s more stable to take the matrix division, where A \ B = inverse(A) * B, but is more stable. So this should be A \ (A_delta - I). But again, don't represent the I` and do this this way. Add the diagonal -1 element as before rather than creating more identity matrices.
Reshape this so it can be just eta[parent[i]].
The to_int can be done in transformed data and saved here. In general, you want to precompute as much as you can in transformed data so that you don’t have to branch elsewhere. For example, all the i such that miss[i, 1] == 0 can be precomputed and stored in an array miss_ys_1 and then the sampling statement is just
if (!prior_only) {
miss_y_1 ~ bernoulli_logit(eta[miss_tip_id_1], 1] + drift_tips[miss_tip_id_1]);
miss_y_2 ~ ...
}
where you’ve also accumulated the indexes for miss_tip_id_1 to match the `miss_y_1.
Can you code prior_only by just setting N_obs = 0?
That was a lot, but I hope you got the general idea of how to think about this. There’s a chapter of the User’s Guide on optimizing code that discusses pretty much everything I suggested in one way or another, at least indirectly.