Function `by`

Modeling
1

I was discussing how to apply a function across groups on Slack and figured to repost on discourse

Let’s say you want to perform some operation across a vector where the elements of the vector each map to a group. It would be really nice if the functions were able to be applied on each group index, which may happen when stan-math upgrades to Eigen 3.4 as @andrjohns noted (see array of integers Eigen).

Let’s say we want to do a dot_product by group. It would look something like this

functions {
vector dot_by(vector x, vector y, array[] int by_index, int J) {
        int N = num_elements(x);
        vector[J] x_dot_by = rep_vector(0., J);

        for (n in 1:N)
            x_dot_by[by_index[n]] += x[n] * y[n];

        return x_dot_by;
  }
}
...
vector[N] x;
vector[N] w;
array[N] int<lower=1, upper=N_groups> groups;

vector[N_groups] = dot_by(x, w, groups, N_groups);

This is not optimal because it does not use dot_product. Once Eigen 3.4 is available this could be written as (I’m not sure this is the exact syntax, @andrjohns may have a better idea, but this is the general idea);

// **WARNING:** below is not valid until Eigen 3.4
vector dot_by(vector x, vector y, array[] int by_index, int J) {
  vector[J] result;
  
for(j in 1:J) 
    result[j] = dot_product(x[by_index], y[by_index]);
  
 return result;
}

We could write the dot_by to work with dot_product in Stan right now ordering the vector by group and using segment. However, this is a bit onerous when one has many groups to do this op over, as each would need to be ordered. Either passing the ordered vectors in as data or doing it in transformed data.

1 Like
2

The general idea would be to have a c++ backend which would take an unsorted array of group indices (e.g., given N individuals and 3 groups, this array would contain N integers in the range 1 - 3). And return a ragged array, containing arrays of the indices for each group.

In c++ pseudo-code, this would look like:

// Integer array containing group membership for 5 individuals
std::vector<int> by_index{1, 1, 3, 1, 2};

// Get number of groups as largest index
int J = *std::max_element(v.begin(), v.end());

// Array to contain arrays of group indices
std::vector<std::vector<int>> group_indices(J);

// Loop over input array of group membership (by_index)
// Assign current index (i) to the matching group index array
// (according to by_index[i])
for (int i = 0; i < N; i++) {
  // Need to subtract 1, as c++ indexing starts at 0
  group_indices[by_index[i] - 1].push_back(i - 1);
}

Once we have that ragged array of group membership (group_indices), we can then loop over it to construct vectorised operations for each group, using the array indexing which will be available in Eigen 3.4.

For example:

// Input vectors of data, of length N
Eigen::VectorXd v1 = ...;
Eigen::VectorXd v2 = ...;

// Vector to hold result for each group
Eigen::VectorXd result(J);

for (int j = 0; j < J; j++) {
  // Compute a grouped dot-product
  result[j] = v1(group_indices[j]) * v2(group_indices[j]).transpose();
}
1 Like
3

Of course, the idea would be that it would be much more generalised by the time it made it to the Stan language, such that any function and number of arguments can be used.

Using the dot_by example above, one way would be something like:

vector[N_groups] = map_by(dot_product, groups, x, w);

Where the map_by function takes the operation (or user-defined function), the array of group indices, and input vector(s) to be indexed. There’s some additional complexity there to be figured out, like working with matrices and including scalars in the function, but overall could be a handy addition

2 Likes