AdvancedPythonProgramming/Chapter07/simul.py at master · Techlord-RCE/AdvancedPythonProgramming · GitHub

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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import numpy as np
from cevolve import c_evolve, c_evolve_openmp

class Particle:

    __slots__ = ('x', 'y', 'ang_speed')

    def __init__(self, x, y, ang_speed):
        self.x = x
        self.y = y
        self.ang_speed = ang_speed

class ParticleSimulator:

    def __init__(self, particles):
        self.particles = particles

    def evolve_cython(self, dt):
        timestep = 0.00001
        nsteps = int(dt/timestep)

        r_i = np.array([[p.x, p.y] for p in self.particles])
        ang_speed_i = np.array([p.ang_speed for p in self.particles])

        c_evolve(r_i, ang_speed_i, timestep, nsteps)

        for i, p in enumerate(self.particles):
            p.x, p.y = r_i[i]

    def evolve_openmp(self, dt):
        timestep = 0.00001
        nsteps = int(dt/timestep)

        r_i = np.array([[p.x, p.y] for p in self.particles])
        ang_speed_i = np.array([p.ang_speed for p in self.particles])

        c_evolve_openmp(r_i, ang_speed_i, timestep, nsteps)

        for i, p in enumerate(self.particles):
            p.x, p.y = r_i[i]

    def evolve_numpy(self, dt):
        timestep = 0.00001
        nsteps = int(dt/timestep)

        r_i = np.array([[p.x, p.y] for p in self.particles])
        ang_speed_i = np.array([p.ang_speed for p in self.particles])

        v_i = np.empty_like(r_i)

        for i in range(nsteps):
            norm_i = np.sqrt((r_i ** 2).sum(axis=1))

            v_i = r_i[:, [1, 0]]
            v_i[:, 0] *= -1
            v_i /= norm_i[:, np.newaxis]

            d_i = timestep * ang_speed_i[:, np.newaxis] * v_i

            r_i += d_i

        for i, p in enumerate(self.particles):
            p.x, p.y = r_i[i]

    def evolve_python(self, dt):
        timestep = 0.00001
        nsteps = int(dt/timestep)

        # First, change the loop order
        for p in self.particles:
            t_x_ang = timestep * p.ang_speed
            for i in range(nsteps):
                norm = (p.x**2 + p.y**2)**0.5
                p.x, p.y = p.x - t_x_ang*p.y/norm, p.y + t_x_ang * p.x/norm

from random import uniform

def benchmark(npart=100, method='python'):
    particles = [Particle(uniform(-1.0, 1.0),
                          uniform(-1.0, 1.0),
                          uniform(-1.0, 1.0))
                  for i in range(npart)]

    simulator = ParticleSimulator(particles)
    if method=='python':
        simulator.evolve_python(1.0)

    if method == 'cython':
        simulator.evolve_cython(10.0)

    if method == 'openmp':
        simulator.evolve_openmp(10.0)

    elif method == 'numpy':
        simulator.evolve_numpy(1.0)