PyScalapack 0.3.6

python wrapper for scalapack

PyScalapack

PyScalapack is a python wrapper for scalapack.

Documents

Load scalapack

It is needed to load scalapack dynamic linked library at first of all step.

from PyScalapack import Scalapack

scalapack = Scalapack("libscalapack.so")

Create a context

Create a blacs context to do other blacs or scalapack operator.

from PyScalapack import Scalapack

scalapack = Scalapack("libscalapack.so")

with scalapack(layout=b'C', nprow=1, npcol=1) as context:
    pass

Create an array

Create a blacs array, and fill it with random generated by numpy.

import numpy as np
from PyScalapack import Scalapack

np.random.seed(0)
scalapack = Scalapack("libscalapack.so")

with scalapack(b'C', 1, 1) as context:
    array = context.array(m=128, n=512, mb=1, nb=1, dtype=np.float64)
    array.data[...] = np.random.randn(*array.data.shape)
    print(array.data)

[[ 1.76405235  0.40015721  0.97873798 ...  1.30142807  0.89526027
   1.37496407]
 [-1.33221165 -1.96862469 -0.66005632 ...  0.70104134 -0.41747735
  -1.09749665]
 [ 1.71230522 -0.79211502 -1.04552456 ...  0.28376955 -0.28219588
  -1.15820319]
 ...
 [-1.47166838  0.82070721 -1.1493715  ...  0.07881221 -2.63213675
   0.75161925]
 [-0.67473808  1.78800397  0.06002943 ... -0.23778156 -1.14289687
  -1.31748978]
 [ 0.26861843  0.26574383 -0.0437187  ... -0.29171979 -0.18048776
   0.37474441]]

Call scalapack function

Call pdgemm and compare it to product calculated in numpy.

import numpy as np
from PyScalapack import Scalapack

np.random.seed(0)
scalapack = Scalapack("libscalapack.so")

with scalapack(b'C', 1, 1) as context:
    array = context.array(m=128, n=512, mb=1, nb=1, dtype=np.float64)
    array.data[...] = np.random.randn(*array.data.shape)

    result = context.array(m=128, n=128, mb=1, nb=1, dtype=np.float64)

    scalapack.pdgemm(
	b'N',
	b'T',
	*(128, 128, 512),
	scalapack.d_one,
	*array.scalapack_params(),
	*array.scalapack_params(),
	scalapack.d_zero,
	*result.scalapack_params(),
    )

    print(result.data - array.data @ array.data.T)

[[0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 ...
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]]

Generic variables

f_one and f_zero is used to get the floating 1 and 0 by selected scalar type, which is useful sometimes.

from PyScalapack import Scalapack

scalapack = Scalapack("libscalapack.so")

print(scalapack.f_one["D"] == scalapack.d_one)
print(scalapack.f_zero["Z"] == scalapack.z_zero)

True
True

Generic functions

Some function such p?gemm could be selected by pgemm[char] where char is one of S, D, C, Z. But this is not applied to all functions, since it is manully mapped. We only map the function we are using currently. If you want to add some other scalapack functions, you could add the mapping by yourself, or just create an issue or pull request.

from PyScalapack import Scalapack

scalapack = Scalapack("libscalapack.so")

print(scalapack.pdgemm)
print(scalapack.pgemm["D"])

<function Scalapack._fortran_function.<locals>.result at 0x7f0a0a35dbd0>
<function Scalapack._fortran_function.<locals>.result at 0x7f0a0a35dbd0>