[Numpy-discussion] Interface for wrapping random distribution functions (__array_function__, __ufunc__, ?)

[Yoann Mocquin]

Hello there,

here is a feature request for the possibility to wrap numpy random distributions to be wrapped by a mechanism like __array_function__ or __array_ufunc__. You can find the GH issue at : https://github.com/numpy/numpy/issues/19382 <https://github.com/numpy/numpy/issues/19382>.



This post follows [this one](https://github.com/numpy/numpy/issues/18902 <https://github.com/numpy/numpy/issues/18902>). I figured I should open another issue for `np.random.normal`, but this applies for all distributions I guess.

## Feature

Basicaly, I would like that `numpy.random.*` distributions could trigger an interface when passed instances from custom classes, "à la" `__array_ufunc__` or `__array_function__`. Here is an example concept : 

```python
arr = np.arange(10)

# Reference result 
print(np.mean(arr))
print(np.random.normal(arr))

custom_obj = MyArrayLike(arr)
print(np.mean(custom_obj))           # OK : np.mean will trigger __array_function__ interface
print(np.random.normal(custom_obj))  # KO : np.random.normal will "only" try to cast the object to float
```

And here is a MWE : 
```python
import numpy as np
np.random.seed(1234)

HANDLED_FUNCTIONS = {}

class NumericalLabeled():
    def __init__(self, value, label=""):
        self.value = value
        self.label = label
        
    def __repr__(self):
        return "NumericalLabelled<"+str(self.value) + "," + self.label+">"
    
    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS:
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)
    

def make_numericallabelled(x, label=""):
    """
    Helper function to cast anything into a NumericalLabelled object.
    """
    if isinstance(x, NumericalLabeled):
        return x
    else:
        return NumericalLabeled(x, label=label)
    
# Numpy functions            
# Override functions - used with __array_function__
def implements(np_function):
    def decorator(func):
        HANDLED_FUNCTIONS[np_function] = func
        return func
    return decorator    
    

@implements(np.random.normal)
def np_random_normal(loc=0.0, scale=1.0, **kwargs):
    # cast both loc and scale into Numericallabelled
    loc = make_numericallabelled(loc)
    scale = make_numericallabelled(scale)
    # check their label is "compatible"
    if not loc.label == scale.label:
        raise ValueError
    return NumericalLabeled(np.random.rand(loc=loc.value,
                                           scale=scale.value, **kwargs), 
                            loc.label+scale.label)

@implements(np.mean)
def np_mean(a, *args, **kwargs):
    return NumericalLabeled(np.mean(a.value, *args, **kwargs),
                            a.label)



def main():
    # reference result for standard array
    arr = np.arange(10)
    print(np.mean(arr))
    print(np.random.normal(arr))
    
    # array-like object
    num_labeled = NumericalLabeled(arr, "toto")
    print(np.mean(num_labeled))
    try:
        print(np.random.normal(num_labeled))
    except Exception as e:
        print(e)

main()
```
which results in 
```
4.5
[ 0.47143516 -0.19097569  3.43270697  2.6873481   3.27941127  5.88716294
  6.85958841  6.3634765   8.01569637  6.75731505]
NumericalLabelled<4.5,toto> 
float() argument must be a string or a number, not 'NumericalLabeled'
```

Since the distribution functions accept array as input, I would expect them to be wrappable like many other array functions.

### Versions
```
Python : 3.8.5 (default, Sep  4 2020, 02:22:02) 
[Clang 10.0.0 ]
Numpy : 1.21.0
```


Cheers
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <https://mail.python.org/pipermail/numpy-discussion/attachments/20210630/de867291/attachment.html>