Python vs Julia

Python vs Julia#

Last week we had a series of lectures from Nils Deppe on partial differential equations. He worked with Python – which makes a nice opportunity to discuss some of the differences between these languages and the strategies that are relevant.

The central difference between Python and Julia is that Python is strictly interpreted, while Julia uses ``just-in-time” compilation. The main upshot is that by default Python is slow. To use Python efficiently for numerical calculations one needs to make use of packages like Numpy and Scipy which call compiled code. For fast operations in Python one always wants to use vectorized code which acts on the vector object in its entirety.

Julia, however, is equally fast using vectorized and element-wise operations. As an example, in Nils Deppe’s code he used the scipy function roll – which cyclically permutes the elements of a vector. To calculate the five-point derivative, Nils used something like

def d_dx_fd4(u, dx):
"""Fourth-order centered first derivative of periodic data u with grid spacing dx.
"""
du_dx = (
    np.roll(u, 2) - 8 * np.roll(u, 1) + 8 * np.roll(u, -1) - np.roll(u, -2)
) / (12 * dx)
return du_dx

There is a Julia command circshift which does the same thing as roll – so we could use that to write the same function.

function vd_dx(u, dx=1)
    du_dx = (
        circshift(u, 2) - 8 * circshift(u, 1) + 8 * circshift(u, -1) - circshift(u, -2)
    ) / (12 * dx)
    return du_dx
end

vd_dx (generic function with 2 methods)

u1=[sin(x) for x in 2pi/100000:2pi/100000:2pi]
@time du1=vd_dx(u1,2pi/100000);

  0.238313 seconds (745.26 k allocations: 44.751 MiB, 25.58% gc time, 99.62% compilation time)

This Julia code should run roughly as fast as the Python code – since they both use a compiled routine (roll or circshift) to do the underlying manipulation.

Where Julia has an advantage is that we can make a more efficient version which makes use of element-wise operations.

function ed_dx(u,dx=1)
    du_dx=similar(u)
    i(j)=mod1(j,length(u))
    for j in eachindex(u)
        du_dx[j]= (u[i(j+2)] -8u[i(j+1)]+8u[i(j-1)]-u[i(j-2)])/(12*dx)
    end
    return du_dx
end

ed_dx (generic function with 2 methods)

u2=[sin(x) for x in 2pi/100000:2pi/100000:2pi]
@time du2=ed_dx(u2,2pi/100000);

  0.014230 seconds (15.73 k allocations: 1.515 MiB, 96.63% compilation time)

Of course – the underlying code is so fast that it doesn’t really matter. The real point is not that the more explicit Julia code is faster – but rather that it is no slower. Sometimes it is more efficient to think element-wise. In Python that will give you a performance hit – while in Julia you can just do it.

There are tricks to make the indexing a little prettier, but

We can possibly make this even more efficient if we give the compiler some extra hints. (The compiler is pretty smart, so these often make no difference @simd @inbounds

function d_dx(u,dx=1)
    du_dx=similar(u)
    i(j)=mod1(j,length(u))
    @inbounds @simd for j in eachindex(u)
        du_dx[j]= (u[i(j+2)] -8u[i(j+1)]+8u[i(j-1)]-u[i(j-2)])/(12*dx)
    end
    return du_dx
end

d_dx (generic function with 2 methods)

u2=[sin(x) for x in 2pi/100000:2pi/100000:2pi]
@time du2=d_dx(u2,2pi/100000);

  0.025004 seconds (32.04 k allocations: 2.260 MiB, 55.67% gc time, 42.55% compilation time)

We got some very small gains from the @inbounds – which turns off certain bounds checking – but not gains from @simd – presumably the compiler already figured out that it could use a simd strategy.

We can get another factor of 3 in speed by separating out the wrapped entries, and saving a lot of mod1 calls

function fd_dx(u,dx=1)
    du_dx=similar(u)
    l=length(u)
    i(j)=mod1(j,l)
    @inbounds @simd for j in 3:l-2
        du_dx[j]= (u[j+2] -8u[j+1]+8u[j-1]-u[j-2])/(12*dx)
    end
    @inbounds @simd for j in [1,2,l-1,l]
        du_dx[j]= (u[i(j+2)] -8u[i(j+1)]+8u[i(j-1)]-u[i(j-2)])/(12*dx)
    end
    return du_dx
end

fd_dx (generic function with 2 methods)

@time fdu2=fd_dx(u2,2pi/100000);

  0.016966 seconds (64.92 k allocations: 3.901 MiB, 99.63% compilation time)

Again – it doesn’t look like the @inbounds or @simd decorators are doing much – and the speed doesn’t really change if one gets rid of them.

Nonetheless, in the end we got about an order of magnitude total speedup over the original code.

Often the other speedup we can get is to pre-allocate the output vector. Here, where we only have a single allocation, it doesn’t really make much of a difference. But the ODE solver can be called with functions like this that put the answer into a set memory location:

function fd_dx!(u,du_dx,dx=1)
    l=length(u)
    i(j)=mod1(j,l)
    @inbounds @simd for j in 3:l-2
        du_dx[j]= (u[j+2] -8u[j+1]+8u[j-1]-u[j-2])/(12*dx)
    end
    @inbounds @simd for j in [1,2,l-1,l]
        du_dx[j]= (u[i(j+2)] -8u[i(j+1)]+8u[i(j-1)]-u[i(j-2)])/(12*dx)
    end
    return du_dx
end

fd_dx! (generic function with 2 methods)

fdu3=similar(u2)
@time fd_dx!(u2,fdu3,2pi/100000);

  0.014764 seconds (55.97 k allocations: 2.718 MiB, 99.67% compilation time)

Views#

The other inefficiency I saw in Nils’ code is that it kept unpacking and repacking the data. At times it wanted to work with a long vector of length 3N, which represented all of the dynamical variables, and other times it wants to break it into three vectors, each of length N. The slick way to deal with this sittuation is to use a View – which is when multiple variables use the same memory address.

function Ψview(v :: Vector)
    l=length(v)
    N=div(l,3)
    @view v[1:N]
end

function Πview(v :: Vector)
    l=length(v)
    N=div(l,3)
    @view v[N+1:2N]
end

function Φview(v :: Vector)
    l=length(v)
    N=div(l,3)
    @view v[2N+1:3N]
end

Φview (generic function with 1 method)

data=collect(1:18)

18-element Vector{Int64}:
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18

Πdata=Πview(data)

6-element view(::Vector{Int64}, 7:12) with eltype Int64:
  7
  8
  9
 10
 11
 12

Πdata[1]=101

Πdata

6-element view(::Vector{Int64}, 7:12) with eltype Int64:
 101
   8
   9
  10
  11
  12

data

18-element Vector{Int64}:
   1
   2
   3
   4
   5
   6
 101
   8
   9
  10
  11
  12
  13
  14
  15
  16
  17
  18

The default reshape function in Julia also supplies a view

ar_data=reshape(data,(6,3))

6×3 Matrix{Int64}:
101  13
  8  14
  9  15
 10  16
 11  17
 12  18

ar_data[2,2]=102
data

18-element Vector{Int64}:
   1
   2
   3
   4
   5
   6
 101
 102
   9
  10
  11
  12
  13
  14
  15
  16
  17
  18

alt_Ψview=@view ar_data[:,1]

6-element view(::Matrix{Int64}, :, 1) with eltype Int64:
 1
 2
 3
 4
 5
 6

alt_Ψview[2]=500
data

18-element Vector{Int64}:
   1
 500
   3
   4
   5
   6
 101
 102
   9
  10
  11
  12
  13
  14
  15
  16
  17
  18

ODE Package#

The last thing to mention is that the ODE package that Nils used in Python is a little more restrictive than the one we used. In particular, he needed to work with a vector input. I believe that the package that we are using can use 2D arrays, or even custom structs. That gives a bit more flexibility.

I won’t go into detail on this – but you can look up the documentation yourself.

Python vs Julia

Contents

Python vs Julia#

Views#

ODE Package#