在numba中实现多态

python

在numba中实现多态

Python是很具存在主义特色的编程语言，萨特在《存在主义是一种人道主义》中写到过，

无作为论是那些说“让别人做我不能做的”的人的态度。我给你们陈述的这种学术恰恰和这种态度相反，因为它宣称除掉行动外，没有真实。确实，它还进一步补充说：“人只是他企图成为的那样，他只是在实现自己意图上方才存在，所以他除掉自己的行动总和外，什么都不是；除掉他的生命外，什么都不是。”正因为如此，所以我们不难理解为什么有些人听到我们的教导感到骇异。因为许多人郁郁不得志时只有一个给自己打气的办法，那就是这样跟自己说：“我这人碰见的事情总是不顺手，否则我的成就要比过去大得多。诚然，我从来没有碰到过一个我真正爱的女人，或者结识过一个真正要好的朋友；不过那是因为我从来没有碰到过一个值得我结识的男人，或者一个真正值得我爱的女人；如果我没有写过什么好书，那是因为我过去抽不出时间来写；还有，如果过去我没有什么心爱的孩子，那是因为我没有能找到可以同我一起生活的男人。所以我的能力、兴趣和能够发挥的潜力，是多方面的，虽然没有用上但是完全可以培养的；因此决不可以仅仅根据我过去做的事情对我进行估价；实际上，我不是一个等闲的人。”但是实际上，而且在存在主义者看来，离开爱的行动是没有爱的；离开了爱的那些表现，是没有爱的潜力的；天才，除掉艺术作品中所表现的之外，是没有的。普鲁斯特的天才就表现在他的全部作品中；拉辛的天才就表现在他的一系列悲剧中，此外什么都没有。为什么我们要说拉辛有能力再写一部悲剧，而这部悲剧恰恰是他没有写的呢？一个人投入生活，给自己画了像，除了这个画像外，什么都没有。当然，这种思想对于那些一生中没有取得成就的人是有点不好受的。另一方面，这却使人人都容易理解到只有实际情况是可靠的；梦、期望、希望只能作为幻灭的梦、夭折的希望、没有实现的期望来解释人；这就是说，只能从反面，而不是从正面来解释。虽说如此，当一个人说，“你除掉你的生活之外，更无别的”，这并不意味着说一个画家只能就他的作品来估计他，因为还有千百件其他的事情同样有助于解释他的为人。这话的意思就是说，一个人不多不少就是他的一系列行径；他是构成这些行径的总和、组织和一套关系。

当看到一只鸟走起来像鸭子、游泳起来像鸭子、叫起来也像鸭子，那么这只鸟就可以被称之为鸭子。在鸭子类型中，我们关注点在于对象的行为，而不是对象的类型。

编程语言可以分为两大类，一类是静态语言，比如C++、Fortran，在使用一个变量之前，必须声明它的类型，而且代码必须通过编译器编译后才能执行。在编译阶段，编译器会检查变量、函数和表达式的类型，确保他们的使用与他们声明的类型相匹配。在C++里，可以通过将基类中的成员函数声明为虚函数来实现运行时多态，也就是说可以通过基类的指针去调用派生类的函数。

下面这个例子中，make_sound函数接受的参数是基类Animal的对象，调用了sound方法，但是因为基类的sound方法被声明为了虚函数，此时当传入Animal的派生类时，make_sound函数会调用派生类重写的方法。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
#include <iostream>
#include <windows.h>

// Animal 类作为基类
class Animal {
public:
    virtual void sound() {
        std::cout << "Unknown animal sound" << std::endl;
    }
};

// Dog 类作为派生类
class Dog : public Animal {
public:
    void sound() override {
        std::cout << "Woof!" << std::endl;
    }
};

// Cat 类作为派生类
class Cat : public Animal {
public:
    void sound() override {
        std::cout << "Meow!" << std::endl;
    }
};

class Animal2 {
public:
    virtual void sound() {
        std::cout << "Unknown animal sound" << std::endl;
    }
};

// make_sound 函数，实现多态并进行类型检查
void make_sound(Animal* animal) {
    animal->sound();
}

int main() {
    Animal animal;
    Animal2 animal2;
    Dog dog;
    Cat cat;

    // 调用 make_sound 函数，实现多态并进行类型检查
    make_sound(&animal); // 输出 "Unknown animal sound"
    make_sound(&dog);    // 输出 "Woof!"
    make_sound(&cat);    // 输出 "Meow!"
    // make_sound(&animal2);    // 编译器报错

    system("pause");
    return 0;
}

如果没有在基类Animal中声明sound为虚函数，那么在调用make_sound函数时，由于sound函数不是虚函数，那么编译器将使用Animal类中的sound函数，而不是调用派生类中的函数。因此此时不会输出"Woof!"或者"Meow!"，只会输出"Unknown animal sound"。其中派生类函数中的override关键字提供了额外的编译时检查。它告诉编译器，派生类中的成员函数必须重写基类中的虚函数，否则无法通过编译。另一方面，即使Animal2类和Animal类实现的一模一样，也没法通过make_sound函数来调用，因为make_sound函数只接受Animal类的对象。

同样的，在下面这个例子中，Jojo类声明了一个基类的属性a，当把这个a赋值为派生类的对象，并调用它的func方法，由于在基类中func声明成了虚函数，这时也会自动地调用派生类的方法。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
#include <iostream>
#include <windows.h>

class Base {
public:
    virtual void func() {
        std::cout << "Base::func() called" << std::endl;
    }
};

class Derived : public Base {
public:
    void func() override {
        std::cout << "Derived::func() called" << std::endl;
    }
};

class Jojo {
public:
    Base* a;
};

int main() {
    Derived derivedObj;
    Jojo jojoObj;
    jojoObj.a = &derivedObj;

    jojoObj.a->func(); // 通过指针调用属性 a 的 func() 方法

    system("pause");
    return 0;
}

在这一点上，对于动态语言Python不需要声明这么多，因为Python不需要在使用一个变量前声明它的类型，不需要进行显式的类型检查，因为 Python 的鸭子类型允许对象的类型通过其行为和属性来确定。因此只要这个对象实现了sound方法，就可以被make_sound函数来调用。我们把这种接口指定的方式，叫作协议（Protocol），只要一个对象具有符合特定协议的方法或属性，就可以被视为符合该协议的对象，不论它的类型是什么，这使得Python的灵活性得以发挥。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
# Animal 类作为基类
class Animal:
    def sound(self):
        print("Unknown animal sound")

# Dog 类作为派生类
class Dog(Animal):
    def sound(self):
        print("Woof!")

# Cat 类作为派生类
class Cat(Animal):
    def sound(self):
        print("Meow!")

class Animal2:
    def sound(self):
        print("Unknown animal sound")

# make_sound 函数，实现多态并进行类型检查
def make_sound(animal):
    animal.sound()

animal = Animal()
animal2 = Animal2()
dog = Dog()
cat = Cat()

# 调用 make_sound 函数，实现多态并进行类型检查
make_sound(animal)  # 输出 "Unknown animal sound"
make_sound(dog)     # 输出 "Woof!"
make_sound(cat)     # 输出 "Meow!"
make_sound(animal2)  # 输出 "Unknown animal sound"

Python中一切皆对象，我们以为自己在操作一个整型，实际上我们只是在操作一个看起来像整型的对象，在C中一个4个字节的整型在Python中占据了28个字节。同时Python并不需要编译之后生成一个可执行文件才能执行，而是Python解释器将源代码解释成字节码，再由虚拟机来执行这些字节码，这个过程带来了极大的性能开销，同样的操作相比于C或Fortran慢两到三个量级是相当正常的。

在numba中实现多态

在Python中直接进行计算密集的任务是不可接受的，如果我们想兼顾Python的灵活性和编译引入静态的高效性，numba是一个折中的选择。numba采用即时编译（Just-In-Time Compilation, JIT Compilation），实时地将Python代码转换为本地的机器代码，这样当函数被调用时，实际上是执行编译后的机器代码，略去了解释过程的开销。可以找到不少介绍numba的文章，一般强调的都是numba非常简单，用一行装饰器就可以直接编译Python原生函数了，而无需改动原本的Python代码。当问题并不复杂，而且不涉及调用其他库的时候，的确是这样的。而一旦问题复杂起来了.. 比如用到了numpy中不太常见的函数，或者试图在numba里写面向对象，这里就会有非常多的限制，最后不得不更改写法来适应numba的需要，倒也没有太多地节省成本，或许直接写C++然后用Pybind11做个包装是更好的选择。。

以numba中装饰类为例，这个时候需要在类上面加的装饰器是

1
@numba.experimental.jitclass(class_spec)

在使用jitclass时，对于类的属性和类中调用的函数，有以下基本的限制。

首先，类的属性必须提前通过一个元组的列表class_spec来声明，每个元组的第一个参数是属性的名字，第二个参数是属性的类型，属性的类型支持Python的原生类型，大部分的numpy数组，以及同样用jitclass类的对象，以及由jitclass类对象构成的列表。

其次，在类中调用函数时，支持主流的Python中的内建函数，如len，部分的numpy函数，其他jitclass对象的函数，以及用jit装饰后的全局函数。也就是说除了Python本身或者numpy的部分函数，剩下能调用的函数必须都经过numba即时编译一遍。函数支持的参数和类属性基本是一致的。

numba支持jitclass去继承若干个非jitclass基类，但是仍需要在class_spec中明确写出所有的属性的名字和类型，而且并不会自动编译这些基类的函数，所以如果想调用这些基类的函数的话，他们也需要满足上述对于调用的函数的要求。

numba不支持jitclass基类的继承和动态派发，也不支持鸭子类型，因为jitclass需要知道类中所有属性的确切类型，这与多态的特性并不兼容。。。

比如考虑这样一个场景，蒙特卡洛模拟，粒子在一个体系中运动，这个结构中可能存在着各种各样的边界，粒子撞到对应的边界了，就要按照对应边界的类型去处理粒子接下来的行为。同时，这个体系中可能有着多个区域，每个区域都有若干个边界组成，这要求我们实现一个区域类，这个类由若干个边界类对象构成。一个基于大量if/else的代码，可能是这么处理不同的边界的：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import numba as nb

Boundary_spec = [
    ("condition", nb.int32),
]

@nb.experimental.jitclass(Boundary_spec)
class Boundary(object):
    def __init__(self, condition):
        self.condition = condition
    
    def collide(self, phonon):
        if self.condition == 0:
            return self.reflect(phonon)
        elif self.condition == 1:
            return self.absorb(phonon)
        # ... and so on for the other conditions
        
    def reflect(self, phonon):
        # implement reflection logic here

    def absorb(self, phonon):
        # implement absorption logic here
        
    # ... and so on for the other behaviors

这种方式在 Boundary 类中为每种可能的边界行为（反射，吸收，等等）都定义一个方法，然后在 collide 方法中，根据 condition 来调用相应的方法。这种实现方式可以利用jitclass来去对整个类去进行即时编译做优化，因为不同的边界行为只是在调用不同的边界方法，所有边界都是Boundary类的对象，这个时候并没有实现多态。

正常情况下在Python中，我们可以通过类的继承或者鸭子类型来实现多态，我们定义一个Boundary基类，然后通过方法重写来实现多态性，为不同的边界类型创建不同的子类，

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Boundary(object):
    def __init__(self, ...):
        ...
        
    def collide(self, phonon):
        raise NotImplementedError()
        

class ReflectBoundary(Boundary):
    def __init__(self, ...):
        super().__init__(...)
        
    def collide(self, phonon):
        ...
        # reflect logic
        ...

class PeriodicBoundary(Boundary):
    def __init__(self, ...):
        super().__init__(...)
        
    def collide(self, phonon):
        ...
        # periodic logic
        ...

# similarly define other boundary classes

这种实现可以用jitclass去包装每一个派生类，但是此时没法把更大的区域类用jitclass去做包装，因为我们不知道这个区域中各个边界到底是什么类型。这时候的思路或许可以是适当地做一些切割，不再去把整个类都去用jitclass做装饰，而是只把用到的属性定义成其他jitclass的成员，计算消耗较大的函数，单独拿出来用jit去做装饰。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
import numba as nb
import numpy as np

Particle_spec = [
    ('x', nb.float64)
]
@nb.experimental.jitclass(Particle_spec)
class Particle(object):
    def __init__(self):
        self.x = 0
        
    def update(self):
        self.x = np.random.random()

Boundary_spec = [
    ("length", nb.float64),
]
class Boundary(object):
    def __init__(self, length):
        self.length = length
        
ScatterBoundary_spec = [
    ("p", nb.float64)
]
@nb.experimental.jitclass(Boundary_spec + ScatterBoundary_spec)
class ScatterBoundary(Boundary):
    __init__Boundary = Boundary.__init__
    def __init__(self, length, p):
        self.__init__Boundary(length)
        self.p = p
    
    def collide(self, particle):
        particle.x = np.random.random()
    
class Domain(object):
    def __init__(self, boundary_list):
        self.boundary_list = boundary_list
        self.particle = Particle()
        
        
domain = Domain([ScatterBoundary(1, 0.5), ScatterBoundary(2, 0.2)])
domain.boundary_list[0].collide(domain.particle)
bound = ScatterBoundary(1, 0.5)
p = Particle()
bound.collide(p)

这时在Domain类里调用还是在外面调用的耗时相差不会太多，希望复杂之后也还差不多是这样。