User defined function (output : vector), how could I assign values in a vectorised way rather than using for loop


#1

Let’s say now I define a function that needs a vector of positions, and mean and standard deviation of a univariate Gaussian distribution, then returns the probability density of Gaussian at those positions as a vector.

Below is a naive way to implement this function (using for loop)

functions{
vector gaussian_pdf(real mu, real sigma, vector z){
vector[num_elements(z)] z_pdf;
for (i in 1:num_elements(z)){
z_pdf[i] = exp( -(z[i] - mu)^2 / (2 * pow(sigma,2) )) * (1/sqrt(2 * pi() * pow(sigma,2)));
}
return z_pdf;
}
}

I was wondering that if there is a way for me to write a more efficient function to achieve the goal?


#2

The square function can work on vectors. There are also a few other ways to simplify the calculations:

z_pdf= exp( -0.5 * square((z- mu) / sigma) )) * (inv_sqrt(2 * pi() * sigma));


#3

Thank you! So what if we have multiple vectors as the input of my function, for example

functions{
vector temp_fun(vector a, vector b){
vector[num_elements(a)] z;
for (i in 1:num_elements(a)){
z[i] = (a[i] * 2) / b[i];
}
return z;
}
}

How can I modify this function to make it more efficient?

Thank you!


#4

vector[num_elements(a)] z = 2 * a ./ b

That one won’t get you as big a speedup, since much of the speed up happens by avoiding adding things to the autodiff stack in broadcasting, but it makes cleaner code.

Also, you can use ``` to quote code and get better formatting on the forum.


#5

Thank you very much. Could you please recommend me some articles or tutorials about speeding up stan program?