JIT and parallelisation¶

This notebook contains some skeletal information about JIT and parallelisation.

[1]:

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import alpro
import time

Just In Time¶

alpro uses just-in-time (JIT) compilation from the numba library. This basically involves putting a lot of @jit decorators around the python functions that do the matrix manipulations.

Here’s an example of how this works with a for loop. Here we make two identical functions but give one a @jit decorator. We call both versions. We see that the first jit call is actually slower than the pure python (because the compilation is deferred until the first call), but for subsequent calls the numba optimised function is way quicker. Just like magic.

[2]:

from numba import jit

@jit(nopython=True)
def sum_with_jit(N):
    my_sum = 0
    for i in range(N):
        my_sum += 1

def sum_without(N):
    my_sum = 0
    for i in range(N):
        my_sum += 1

t1 = time.time()
x = sum_without(100000)
print ("time, pure python:", time.time() - t1)

t1 = time.time()
x = sum_with_jit(100000)
print ("time, jit (call 1):", time.time() - t1)

t1 = time.time()
x = sum_with_jit(100000)
print ("time, jit (call 2):", time.time() - t1)

time, pure python: 0.004478931427001953
time, jit (call 1): 0.09760212898254395
time, jit (call 2): 5.626678466796875e-05

Exactly the same thing happens with alpro. If we only do one calculation, the calculation is quite slow, but subsequent calculations are nice and fast.

So, be aware of this when you are writing code to use alpro!

[3]:

import alpro
s = alpro.Survival("1275b")
s.init_model()
s.set_params(1e-12 * 1e-9, 1e-13)
s.domain.create_box_array(1800, 0, s.coherence_func, r0=0)
energies = np.logspace(3,4,1000)
t1 = time.time()

for n in range(2):
    t1 = time.time()
    P, Pradial = s.propagate(s.domain, energies, pol="both")
    print ("time:", time.time() - t1)

time: 2.9819839000701904
time: 0.11284494400024414

Parallelisation¶

If you have installed alpro with mpi4py installed then you will have access to a few routines to. The function in question is parallel.run, which takes the following arguments:

function    callable
            function to compute

iterable    iterable
            iterable containing first arguments

kwargs      dictionary
            dictionary of keyword arguments to pass to function

split       str
            how to split up the parallel processes

The function must take as its first argument a variable that is consistent with the type of iterable. For example, if your function takes a string as an argument, then an appropriate iterable would be a list of strings. In alpro, I normally use it for iterating over the RNG seed. An example is shown below, which also shows you how to use the kwargs parameter.

[4]:

def my_calculation(i, g = 1e-12 * 1e-9):
    s = alpro.Survival("1275b")
    s.init_model()
    s.set_params(g, 1e-13)
    s.domain.create_box_array(1800, 0, s.coherence_func, r0=0)
    s.propagate(s.domain, energies, pol="both")

import alpro.parallel as parallel
iterable = range(0,10)
kwargs = {"g": 1e-1 * 1e-9}
time_taken = parallel.run(my_calculation, iterable, kwargs = kwargs)

This is thread 0 calculating models 0 to 9, total 10
Waiting for other threads to finish...
Thread 0 took 1.1756579875946045 seconds to calculate 10 models

This is currently executing in serial, but to run in parallel you could save the above code as script.py and run with, e.g.

mpirun -n 4 python script.py

JIT and parallelisation

Contents

JIT and parallelisation¶

Just In Time¶

Parallelisation¶