Pair-wise distances

Python can be easily interfaced with other programming languages. Here we compute the euclidean distance matrix with componenents $r_{ij} = \left\| \mathbf{ r_i} - \mathbf{ r_j} \right\|$ using FORTRAN. It is possible to seamlessly compile and run FORTRAN code within a Jupyter notebook (similarly to Cython or Pythran) using a magic command, namely %load_ext fortranmagic. This requires the installation of the fortran-magic library. In a command line interface type pip install fortran-magic

Load external libraries.

In [1]:
import math
import numpy as np
import numexpr as ne
import numba as nb
import matplotlib.pyplot as plt
from scipy.spatial.distance import pdist
import os
n_cpu = os.cpu_count()
In [2]:
# Atomic coordinates within cluster
coords = np.loadtxt("Au.xyz")
N, dim = np.shape(coords) #Extract the number of atoms and the number of dimensions

Cython implementation

In [3]:
%load_ext Cython
In [4]:
%%cython --compile-args=-fopenmp --link-args=-fopenmp
import numpy as np
from cython.parallel import prange
import cython

from libc.math cimport sqrt

@cython.wraparound(False)
@cython.boundscheck(False)
def r_ij_cython(double[:,::1] coords):
    cdef:
        double[::1] r
        double tmp
        int i,j,k,l
        int N = coords.shape[0]
        int dim = coords.shape[1]
    
    r = np.zeros(int((N*N-N)/2), dtype=np.float64)
    for i in prange(N, nogil=True):
        for j in range(i+1,N):
            l = int(i * (N - 1) - i * (i + 1) / 2 + j - 1)
            tmp = 0.0
            for k in range(3):
                tmp = tmp + (coords[i,k]-coords[j,k])**2
            r[l] = sqrt(tmp)
    return r
In [5]:
%timeit r_ij_cython(coords)
r_ij_ct = np.asarray(r_ij_cython(coords))
100 loops, best of 3: 11.9 ms per loop

Fortran implementation

In [6]:
%load_ext fortranmagic
In [7]:
%%fortran --f90flags='-fopenmp' --extra='-lgomp'
subroutine getrij(coords,r, N, M)

    use iso_c_binding, only: c_double, c_int
    use omp_lib
    implicit none
    
    real*8, dimension(N*3), intent(in)    :: coords
    real*8, dimension(M),   intent(inout) :: r
    integer(c_int),         intent(in)    :: N
    integer(c_int),         intent(in)   :: M

    integer :: i,j,k,l,ii,jj
    real(c_double) :: tmp
    
    l = 0
    !$OMP PARALLEL DO
    do i = 0 , N-1
        ii = i * 3
        do j = i + 1 , N-1
            jj = j * 3
            tmp = 0.0
            do k = 1 , 3
                tmp = tmp + (coords(ii+k) - coords(jj+k)) ** 2
            end do
            l = 1 + (i * (N - 1) - i * (i + 1) / 2 + j - 1)
            r(l) = sqrt(tmp)
        end do
    end do

end subroutine getrij
In [8]:
def r_ij_fortran(coords):
    N = coords.shape[0]
    M = (N * N - N) // 2
    r = np.zeros(M, dtype=np.float64, order='F')
    getrij(coords.ravel(),  r, N, M)
    return r
  
print("Is the result correct:",np.allclose(r_ij_fortran(coords), r_ij_ct))
%timeit r_ij_fortran(coords)
Is the result correct: True
100 loops, best of 3: 10.2 ms per loop