Is it possible to numpy.vectorize an instance method?

允我心安 提交于 2019-12-01 16:44:21

Simple solution without modifying the class

You can use np.vectorize directly on the method on the instance:

class Dummy(object):

    def __init__(self, val=1):
        self.val = val

    def f(self, x):
        if x == 0:
            return self.val
        else:
            return 2


vec_f = np.vectorize(Dummy().f) 


def test_3():
    assert list(vec_f([0, 1, 2])) == [1, 2, 2]

test_3()

You can also create a vectorized function vec_f in your __init__:

Adding a vectorized version to the instance

class Dummy(object):

    def __init__(self, val=1):
        self.val = val
        self.vec_f = np.vectorize(self.f) 

    def f(self, x):
        if x == 0:
            return self.val
        else:
            return 2


def test_3():
    assert list(Dummy().vec_f([0, 1, 2])) == [1, 2, 2]

or with a different naming scheme:

class Dummy(object):

    def __init__(self, val=1):
        self.val = val
        self.f = np.vectorize(self.scalar_f) 

    def scalar_f(self, x):
        if x == 0:
            return self.val
        else:
            return 2


def test_3():
    assert list(Dummy().f([0, 1, 2])) == [1, 2, 2]

test_3()

    test_3()

Remembering a technique I saw in the memoized decorator, I managed to get the decorator to also work for instance methods by subclassing numpy.vectorize as follows:

import numpy as np
import functools


class vectorize(np.vectorize):
    def __get__(self, obj, objtype):
        return functools.partial(self.__call__, obj)

Now if I decorate the Dummy class' f method with vectorize instead of np.vectorize, the test passes:

class Dummy(object):
    def __init__(self, val=1):
        self.val = val

    @vectorize
    def f(self, x):
        if x == 0:
            return self.val
        else:
            return 2


def test_3():
    assert list(Dummy().f([0, 1, 2])) == [1, 2, 2]

if __name__ == "__main__":
    pytest.main([__file__])

with output

test_numpy_vectorize.py .

=========================== 1 passed in 0.01 seconds ===========================
[Finished in 0.7s]

Here's a generic decorator that works with instance methods as well as functions (refer to Numpy's documentation for otypes and signature):

from functools import wraps

import numpy as np

def vectorize(otypes=None, signature=None):
    """Numpy vectorization wrapper that works with instance methods."""
    def decorator(fn):
        vectorized = np.vectorize(fn, otypes=otypes, signature=signature)
        @wraps(fn)
        def wrapper(*args):
            return vectorized(*args)
        return wrapper
    return decorator

You may use it to vectorize your method as follows:

class Dummy(object):
    def __init__(self, val=1):
        self.val = val

    @vectorize(signature="(),()->()")
    def f(self, x):
        if x == 0:
            return self.val
        else:
            return 2


def test_3():
    assert list(Dummy().f([0, 1, 2])) == [1, 2, 2]

The key is to make use of the signature kwarg. Parenthesized values to the left of -> specify input parameters and values to the right specify output values. () represents a scalar (0-dimensional vector); (n) represents a 1-dimensional vector; (m,n) represents a 2-dimensional vector; (m,n,p) represents a 3-dimensional vector; etc. Here, signature="(),()->()" specifies to Numpy that the first parameter (self) is a scalar, the second (x) is also a scalar, and the method returns a scalar (either self.val or 2, depending on x).

$ pytest /tmp/instance_vectorize.py
======================= test session starts ========================
platform linux -- Python 3.6.5, pytest-3.5.1, py-1.5.3, pluggy-0.6.0
rootdir: /tmp, inifile:
collected 1 item

../../tmp/instance_vectorize.py .                                                                                                                                                     [100%]

==================== 1 passed in 0.08 seconds ======================

If you want to use vectorized implementation of your method you can use excluded parameter like following:

class MyClass:
    def __init__(self, data):
        self.data = data
        self.my_vectorized_func = np.vectorize(self.my_func, excluded='self')

    def my_func(self, x):
        return pow(x, self.data)

With this, you can use your method like the non-vectorized one:

 In[1]: myclass = MyClass(3) # '3' will be the power factor of our function
 In[2]: myclass.my_vectorized_func([1, 2, 3, 4, 5])
Out[3]: array([  1,   8,  27,  64, 125])

From the docs:

The data type of the output of vectorized is determined by calling the function with the first element of the input. This can be avoided by specifying the otypes argument.

The first input in your function f(self, x) is self. Maybe you can make that function a wrapper around a staticmethod function?

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!