[FEATURE] `Adapter.gather()`

[jerrymhuang]

Feature Description
Sometimes, users may want to cherry-pick only a portion of the simulation outputs under a set of customized indices. Take, for example:

# Here is a simple 5x5 matrix
mat = np.arange(25).reshape((5, 5)) 

# Indices for the lower triangle portion of mat (15 indices, in this case)
ix, iy = np.tril_indices(5) 

# A flattened vector of shape (15,) extracted from mat
picks = mat[ix, iy] 

This can happen during data configuration. The Adapter could be used to take care of some simple scenarios like this by allowing the users to specify the portions they would like to pick() from the simulation output.

Example Usages

1. (simple scenario)

sim = {'x': np.random.rand(25)}
picks = [1, 2, 3, 5, 8]

adapter = (
    bf.adapters.Adapter()
    # ...
    .pick('x', indices=picks)
    # ...
)

2 (this scenario follows the example provided at the beginning)

sim = {'a': np.random.rand(25).reshape((5, 5))
ix, iy = np.tril_indices(5)

adapter = (
    bf.adapters.Adapter()
    # ...
    .pick('a', indices=[ix, iy])
    # ...
)

[vpratz]

I see the use of this, and maybe we can find a better name for it than pick. As far as I can tell, the simple scenario is already covered by the Take transform. The more complex case corresponds to gather_nd in TensorFlow and gather in Jax, or simply to slicing in numpy. I'd suggest we use something related to gather or slice, depending on the exact implementation you have in mind. What do you think?

[jerrymhuang]

The name pick was intended to be a placeholder. We could, ideally, extend the existing functionality of take to incorporate the more complex scenarios. I think this makes the most sense since the simple case is already covered.

Otherwise, if this were to be a standalone, then I think gather would be a good name.