# Using map_rect for operations on each column of a matrix

Dear all,

I am quite new in Stan. I want to apply a function (that calculates one element of a vector using the value of the previous one) on each column of a matrix. Instead of doing the calculation one column after another, I am trying to implement map_rect where the shards are the columns of the matrix. Bellow is my Stan code with and without using map_rect.

Without map_rect :

``````functions {
vector function1( real theta, vector x )
{
vector[rows(x)] x_transformed;
x_transformed[1] = x[1];
for ( i in 2:rows(x) ) {
x_transformed[i] = x[i] + theta * x_transformed[i-1];
}
return x_transformed;
}
matrix function1Mat( vector theta, matrix x )
{
int N = rows(x);
int K = cols(x);
matrix[N, K] out;
for ( i in 1:K ) {
out[ 1:N, i ] = function1( theta[i], x[1:N, i] );
}
return out;
}
}
data {
int<lower = 0> N;
int<lower = 0> K;
matrix[N, K] X;
vector[N] Y;
}
transformed data {
}
parameters {
vector[K] beta;
real<lower=0> sigma;
vector<lower=0, upper=1>[K] theta;
real<lower=0> alpha;
}
model {
beta ~ normal(1, 1);
theta ~ normal(0.5, 0.5);
sigma ~ normal(0, 200);
alpha ~ normal(100, 50);
Y ~ normal(alpha + function1Mat(theta, X) * beta, sigma);
}

``````

With map_rect :

``````functions {
vector function1(vector phi, vector theta, real[] x_r, int[] x_i) {
vector[size(x_r)] x_transformed;
x_transformed[1] = x_r[1];
for ( i in 2:size(x_r) ) {
x_transformed[i] = x_r[i] + theta[1] * x_transformed[i-1];
}
return x_transformed;
}
}
data {
int<lower = 0> N;  // rows
int<lower = 0> K;  // cols
matrix[N, K] X;
vector[N] Y;
}
transformed data {
real x_rs[K, N] = to_array_2d(X');  // the shards are the columns of the input matrix
int x_is[K, N];  // x_is is not used
for (k in 1:K) {
x_is[k, 1] = N;
x_is[k, 2] = K;
}
}
parameters {
vector[K] beta;
real<lower=0> sigma;
vector<lower=0, upper=1>[1] theta[K];
real<lower=0> alpha;
vector[0] phi;  // phi is not used
}
model {
alpha ~ normal(100, 50);
beta ~ normal(1, 1);
for (i in 1:K) {
theta[i] ~ normal(0.5, 0.5);
}
sigma ~ normal(0, 200);
Y ~ normal( alpha + to_matrix(map_rect(function1, phi, theta, x_rs, x_is), N, K) * beta, sigma);
}

``````

However, the calculation without map_rect is way faster than with map_rectâ€¦
Without map_rect : â€śChain 1: Gradient evaluation took 0.000199 secondsâ€ť
With map_rect : â€śChain 1: Gradient evaluation took 0.024525 secondsâ€ť

Is there a way to modify the code so that map_rect is effective to speed up the calculation in that case?