Consider simplification or replacement of BlockArray

[Michael-T-McCann]

BlockArrays are/have been at the center of several issues (#173 #159 #149 #147 #146), are difficult to maintain (because they touch jax internals), and necessitate a large amount of code bloat in SCICO (scico.numpy, scico.scipy).

Let's use this issue to discuss possible replacements BlockArray. What functionality would this replacement need to provide? How could we go about it (implementation-wise)?

[bwohlberg]

Current usage examples include:

1. BlockArray is integral part of the operator.BiConlvolve interface
2. BlockArray allows linop.FiniteDifference to be simultaneously applied to multiple axes when the output shape is not the same for all axes
3. We do elementwise arithmetic on BlockArray, e.g.,

   scico/scico/solver.py

   Line 287 in c067188

   r = b - Ax

4. We compute norms on BlockArray, e.g.,

   scico/scico/solver.py

   Line 290 in c067188

   num = r.ravel().conj().T @ z.ravel()

5. We make BlockArrays filled with zeros, e.g., snp.zeros(((1, 2), (3, 4)))

[Michael-T-McCann]

Voting to close this. @crstngc

After consideration (and trying to remove scico.numpy in a branch) I am convinced (1) whatever BlockArray is, we will want helper functions/wrappers to make use of them. Prewrapping numpy in scico.numpy is good solution to this. (2) If we want to support x + y for BlockArrays, then they've got to be a class. (1) and (2) together make the current implementation look like a good solution.

The only modification that might make sense is to change the internal representation to a Tuple[DeviceArray] (i.e. derive from tuple), which (I think) would allow arrays on different devices and avoid some copying. It may also allow simplifications in the implementation (e.g., indexing) that would reduce coupling to jax internals.

Notes on a few options that inform the above reasoning:

1. BlockArray goes away, all functionality comes from jax's pytrees, e.g., Tuple[DeviceArray]. Pros: no magic, no multi-page docs, no confusion about how broadcasting should work. If you know tuples, lists, dicts, you know how these work. These should automatically work with all of jax. Cons: all math becomes cumbersome (probably would want helper functions):
   - to add two together: jax.tree_map(operator.add, x, y).
   - call unary function: jax.tree_map(jnp.sqrt, x) or a wrapped function scico.sqrt(x)
   - reductions: jax.tree_util.tree_reduce or a wrapped function
   - creating one full of zeros: ugly but possible using tree_map's is_leaf argument.
2. BlockArray is a class, but scico.numpy goes away. Pros: convenient syntax for some operators, users only have jnp and np to get confused about. Cons: some math remains cumbersome unless we add helper functions.
   - to add two together: x + y
   - other things, same as (1)
3. (current status) BlockArray is a class, all/most of numpy is prewrapped to work in scico.numpy. Pros: already implemented, convenient from a user perspective. Cons: mentioned in the issue. Also, current implementation could possibly be improved (from contiguous array to Tuple[DeviceArray] underneath?).

[crstngc]

This is a good summary of where we are and I agree that we can close the issue. I will explore the change to an internal representation based on Tuple[DeviceArray] and we can go from there.