Add boolean-based indexing (e.g., Y[Z == 1])

[andrjohns]

Summary:

A useful addition for simpler expressions with more complex indexing/subsetting would be to support logical indices. On the backend, this could just be syntactic sugar which first constructs an array of indices of the true values before passing to the existing indexing implementations.

C++ example:

template <typename EigVec>
inline auto rvalue(EigVec&& v, const char* name, std::vector<bool> lgl_idx) {
  std::vector<int> idx;
  for (int i = 0; i < lgl_idx.size(); i++) {
    if (lgl_idx[i]) {
      idx.push_back(i);
    }   
  }
  return rvalue(std::forward<EigVec>(v), name, index_multi(idx));
}

A simpler alternative, we could add an std::vector<bool> constructor for index_multi which performs the same logic above.

This also has a bit of overlap with the use of bool-ish types in the Stan languages, so there might be some edge-case/interaction that I'm not thinking of. Thoughts?

Current Version:

v2.35.0

[bob-carpenter]

Yes, please. I've wanted this for a long time. Working from Stan itself, I think we'd like to be able to do this:

vector[N] y = ...;
...
array[N] int<lower=0, upper=1> bool_idx = y > 0;
y[bool_idx]

This would require two things:

1. Extending y > 0 and other conditionals to apply to arrays/vectors y to return an array of 0/1 integer values of the same size. Here this requires broadcasting the 0.
2. Allow y[bool_idx] style indexing.

In C++, bool is just shorthand for the integer type that takes values 0 or 1.

There are many more things like this that would be useful. It'd be useful to do a survey of operations provided by the tidyverse and numpy/pandas to see what'd be useful---I often feel while writing Stan code that I want more of these operators and wind up having to write a lot of loops.

[andrjohns]

I don't think we'd be able to treat an array of 0/1 ints as booleans for indexing here unfortunately, since it would be ambiguous whether an array of 1 should be treated as "all TRUE, return whole vector" or "broadcast the first element of vector"

[bob-carpenter]

That's a good point. But we don't have a specific boolean type in Stan, do we? And I think in C++ it'd wind up being ambiguous because a bool is an int. One thing we could do is introduce another notation, like Ts select(Ts x, int[] idxs). I might have said x@idxs, but we're using @ for annotations. It would be nice to have general enough comprehensions we could do something like x[[n in 1:N s.t. idxs[n]] but that's really clunky and involves a lot of implicit binding in the parsing.

The other place I want something like select is when I have

matrix[M, N] x;
array[J] int<lower=1, upper=M> row_idxs;
int<lower=1, upper=N>[J] col_idxs;
vector[J] y = pairwise_index(x, row_idxs, col_idxs);

where pairwise_index(x, row_idxs, col_idxs)[j] = x[row_idxs[j], col_idxs[j]]. I need that function all the time. Or maybe it'd be better with tuple indexes as follows.

matrix[M, N] x;
array[J] tuple(int<lower=1, upper=M>, int<lower=1, upper=N>) idxs;
vector[J] = x[idxs];

[jachymb]

Is adding a native stan boolean type up to discussion? That would allow implementing this by simply adding the appropriate overload to operator[] and modifying the comparison operators to return booleans. This is obviously a breaking change, but maybe it could be considered for a future major release.

[bob-carpenter]

Is adding a native stan boolean type up to discussion?

Everything's fair for discussion.

This is obviously a breaking change,

I would personally be opposed to breaking backward compatibility for the sake of adding a proper boolean type. I would prefer to just add helper functions to do the indexing as I suggested above. I'm also not sure how the type of overload you're suggesting would play out in the C++ implementation, where there's a lot of fluidity between numerical and boolean types.

[andrjohns]

There's some previous discussion around a bool type here: stan-dev/stanc3#1328, essentially treating the type as syntatic sugar for int<lower=0, upper=1> for backwards compatibility

[bob-carpenter]

Thanks for the pointer, @andrjohns. If we took the overloading route, which would be straightforward, it wouldn't allow us to support boolean-based indexing because applying an array to an array of int is already defined. That's why I've been thinking separate function like select(x, idxs) that would be equal to x[idxs] in languages that let you select with arrays of booleans.