Parallel Computing in QuSpin
This example shows how to speed up QuSpin code via multi-threading by using
OpenMP,
Intel’s MKL library for NumPy/SciPy (which is Anaconda’s default NumPy version, starting from Anaconda 2.5 onwards).
Quspin has OpenMP support built in starting from version 1.0.0 (see also Parallel computing support); to install it, simply run
$ pip install quspin
The example below demonstrates how to use the OpenMP version of quspin for parallel computing. It is set up in such a way that the number of OpenMP and MKL threads is controlled from the command line [cf. code lines 8,9]. To run the script, run
$ python example12.py ${OMP_NUM_THREADS} ${MKL_NUM_THREADS}
You can directly compare the speed for different values of the number of threads [make sure your machine’s processor has more than one core]
$ python example12.py 1 1 # single-threaded computation
$ python example12.py 4 1 # multi-threaded OpenMP computation, speedup for basis functions, evolution, and matrix-vector multiplication
$ python example12.py 1 4 # multi-threaded MKL computation, speedup for diagonalization-like routines
$ python example12.py 4 4 # simulaneous OpenMP and MKL multi-threading speedup
Notice how, as explained in Parallel computing support, OMP_NUM_THREADS improves the speed of the basis computation (code line 43), and the time evolution (code line 65), while MKL_NUM_THREADS improves the speed of the exact diagonalization step (code line 60).
Note: there is a common problem with using OpenMP on OSX with anaconda packages for Python 3, which may induce an error unrelated to QuSpin:
$ OMP: Error #15: Initializing libiomp5.dylib, but found libomp.dylib already initialized.
However, this error can be disabled [at one’s own risk!] until it is officially fixed, see code line 7 below.
Script
1###########################################################################
2# example 12 #
3# In this script we show how to use QuSpin's OpenMP and MKL capabilities. #
4###########################################################################
5from quspin.operators import hamiltonian
6from quspin.basis import spin_basis_general
7from quspin.operators._make_hamiltonian import _consolidate_static
8import numpy as np
9from scipy.special import comb
10
11from time import time # timing package
12import os
13import sys
14
15#
16os.environ["KMP_DUPLICATE_LIB_OK"] = (
17 "True" # uncomment this line if omp error occurs on OSX for python 3
18)
19os.environ["OMP_NUM_THREADS"] = str(
20 int(sys.argv[1])
21) # set number of OpenMP threads to run in parallel
22os.environ["MKL_NUM_THREADS"] = str(
23 int(sys.argv[2])
24) # set number of MKL threads to run in parallel
25#
26np.random.seed(1) # fixes seed of rng
27
28
29def run_computation():
30 #
31 ###### define model parameters ######
32 J1 = 1.0 # spin=spin interaction
33 J2 = 0.5 # magnetic field strength
34 Omega = 8.0 # drive frequency
35 Lx, Ly = 4, 4 # linear dimension of spin-1/2 2d lattice
36 N_2d = Lx * Ly # number of sites for spin-1/2
37 #
38 ###### setting up user-defined symmetry transformations for 2d lattice ######
39 sites = np.arange(N_2d) # sites [0,1,2,....]
40 x = sites % Lx # x positions for sites
41 y = sites // Lx # y positions for sites
42 #
43 T_x = (x + 1) % Lx + Lx * y # translation along x-direction
44 T_y = x + Lx * ((y + 1) % Ly) # translation along y-direction
45 #
46 T_a = (x + 1) % Lx + Lx * ((y + 1) % Ly) # translation along anti-diagonal
47 T_d = (x - 1) % Lx + Lx * ((y + 1) % Ly) # translation along diagonal
48 #
49 ###### setting up bases ######
50 basis_2d = spin_basis_general(
51 N_2d, pauli=False
52 ) # making the basis: sped up by OpenMP if symmetries are on
53 print("finished computing basis")
54
55 #
56 ###### setting up hamiltonian ######
57 # set up time-dependence
58 def drive(t, Omega):
59 return np.cos(Omega * t)
60
61 drive_args = [
62 Omega,
63 ]
64 # setting up site-coupling lists
65 J1_list = [[J1, i, T_x[i]] for i in range(N_2d)] + [
66 [J1, i, T_y[i]] for i in range(N_2d)
67 ]
68 J2_list = [[J2, i, T_d[i]] for i in range(N_2d)] + [
69 [J2, i, T_a[i]] for i in range(N_2d)
70 ]
71 #
72 static = [["xx", J1_list], ["yy", J1_list], ["zz", J1_list]]
73 dynamic = [
74 ["xx", J2_list, drive, drive_args],
75 ["yy", J2_list, drive, drive_args],
76 ["zz", J2_list, drive, drive_args],
77 ]
78 # build hamiltonian
79 H = hamiltonian(
80 static,
81 dynamic,
82 basis=basis_2d,
83 dtype=np.float64,
84 check_symm=False,
85 check_herm=False,
86 )
87 # diagonalise H
88 E, V = H.eigsh(time=0.0, k=50, which="LA") # H.eigsh sped up by MKL
89 print("finished computing energies")
90 psi_0 = V[:, 0]
91 # evolve state
92 t = np.linspace(0.0, 20 * 2 * np.pi / Omega, 21)
93 psi_t = H.evolve(psi_0, t[0], t, iterate=True) # H.evolve sped up by OpenMP
94 for j, psi in enumerate(psi_t):
95 E_t = H.expt_value(psi, time=t[j])
96 print("finished evolving up to time step {:d}".format(j))
97
98
99# time computation
100ti = time() # start timer
101run_computation()
102print("simulation took {0:.4f} sec".format(time() - ti))