feat(mpiarray): ufunc, __array_finalize__, and __getitem__ handling for MPIArray

[anjakefala]

If used with draco, requires: radiocosmology/draco#125

Subclassing NumPy arrays 101

view(MPIArray) -> __array_finalize__
MPIArray[slice] -> __getitem__ -> __array_finalize__
MPIArray() -> __new__ -> __array_finalize__

ufuncs are the universal functions that are applied element by element in nparrays. Such a np.add(), np.multiply(). If they end up summing over an axis, they are a ufunc with a reduce method. If they go element-by-element, they are a ufunc with an outer method. If they occur in-place, they have an at method.

ufunc -> __array_ufunc__

More links on the role all these various functions play in writing subclasses for NumPy arrays can be found here: https://github.com/chime-experiment/Pipeline/issues/81

New Exceptions

The new AxisException will be added. It will be raised when there are issues involving the integrity of the distributed axis with MPIArrays.

getitem

• When you use global_slice to index into a distributed axis, then it will return an array for the rank on which the index exists, and None otherwise.
• When you slice on a non-distributed axis, __getitem__ will return an MPIArray, whose distributed axes number might be lower than the original array, depending on if the slice results in an axis reduction. It is assumed that the distributed axis length is unchanged. An mpiutil.split_local call will be made.
• When you directly slice into a distributed axis, it will index into the local arrays on each rank, and return a regular numpy array. This behaviour will be deprecated, and replaced with the raising of an AxisException.

ufunc

• Upon a ufunc being executed, both inputs and outputs are converted to
  standard nparrays, and then the nparray ufunc is called
• all inputted MPIArrays must be distributed along the same axes. An AxisException will be called, if they are not.
• if no output are provided, the results are converted back to MPIArrays. The new array will either be distributed over that axis, or possibly one axis down for reduce methods. keepdims kwargs are handled.
• If you pass a kwarg axis to the ufunc, it must not be the distributed axis. Operations along the distributed axis are not allowed.
• For operations that normally return a scalar, the scalars will be wrapped into a 1D array, distributed across axis 0.
• To produce the MPIArray output, MPIArray.wrap(..) is called. This means there is a call to mpiutil.split_local and mpiutil.allreduce.

array_finalize

This is called whenever a user uses .view(), __new__, and broadcasts on an MPIArray. It finalizes the creation of the output MPIArray. view()s can occur with __getitem__ calls.

If it is a __new__ call, it does nothing.

If we are in an np.ndarray.view() call, it does nothing. This should only occur when we are within a wrap()

If we are in a view(), it grabs the attributes from the origin array.

Misc

If a user wishes to create an MPIArray from an ndarray, they should use MPIArray.wrap(). They should not use ndarray.view(MPIArray).

https://github.com/chime-experiment/Pipeline/issues/81

[jrs65]

@anjakefala we touched on this yesterday, but what I think would help a lot would be to collect a list of MPIArray calls that we would like to work and agree on what their behaviour should be (for which we might want to circulate around to chime-analysis for opinions). Then we can turn them all into unit test cases.

Something like:

# Setup array
comm = MPI.COMM_WORLD
dist_array = mpiarray.MPIArray((comm.size, 4), comm=comm, axis=0)
dist_array[:] = comm.rank

# I think this one is uncontroversial
assert dist_array.sum(axis=1) == 4 * comm.rank

# Should this be allowed?
# assert dist_array.sum(axis=0) == ???

# What should this do? It seems the two sensible options are that is sums over all axes, giving the same scalar everywhere...
assert dist_array.sum() == 4 * comm.size * (comm.size - 1) // 2

# ... or that it just ignores the distributed axis, giving a distributed 1-d array
assert dist_array.sum() == 4 * comm.rank

[lgtm-com]

This pull request introduces 1 alert when merging d2c1a05 into 4a85a3d - view on LGTM.com

new alerts:

• 1 for Suspicious unused loop iteration variable

[jrs65]

Also, note that .all() in an MPIArray reduction over all axes, and we hadn't clearly worked out the semantics for that yet.

Seeing how it works with .all() maybe pushes it in the direction of you should get one value per rank.

[anjakefala]

Linter is failing due to pylint updates. This PR is ready for another conversation now.

[anjakefala]

😭