TDSE with MPI

Here you will learn how to speed up you TDSE code using MPI. First a new-format Richmol file, containing energies of field-free states and matrix elements of field tensor operators for water, is downloaded from zenodo:

import urllib.request
import os
import h5py

# GET RICHMOL FILE
richmol_file = "h2o_p48_j40_emax4000_rovib.h5"
if not os.path.exists(richmol_file):
    url = "https://zenodo.org/record/4986069/files/h2o_p48_j40_emax4000_rovib.h5"
    print(f"download richmol file from {url}")
    urllib.request.urlretrieve(url, "h2o_p48_j40_emax4000_rovib.h5")
    print("download complete")

Time-propagation without MPI

Consider the interaction between electric dipole moment and external dc field. The needed cartesian tensor operators and external field are initialized:

from richmol.field import CarTens
from richmol.convert_units import Debye_x_Vm_to_invcm

h0 = CarTens(richmol_file, name='h0')
mu = CarTens(richmol_file, name='dipole') \
    * (-0.5 * Debye_x_Vm_to_invcm())

# INITIALIZE EXTERNAL DC FIELD
field = lambda t : [0, 0, t * 1e5 ] # (V/m)

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

A TDSE solver and initial state vectors, spanning the space of field-free states, are initialized:

from richmol.tdse import TDSE

# INITIALIZE TDSE AND STATES
tdse = TDSE(t_start=0, t_end=1, dt=0.01, t_units="ps", enr_units="invcm")
vecs = tdse.init_state(h0, temp=None)
print("single process propagates {} vecs".format(len(vecs)))

single process propagates 8995 vecs

The initial states are propagated in time:

import time

time_0 = time.time()
for i, t in enumerate(tdse.time_grid()):
    vecs_, _ = tdse.update(
        mu.field(field(t)), H0=h0, vecs=vecs, matvec_lib='scipy'
    )
    vecs = vecs_
runtime = time.time() - time_0
print('runtime without MPI : {} sec'.format(round(runtime, 3)))

runtime without MPI : 79.436 sec

Time-propagation with MPI

The next code block can be ignored, since it is only needed to intialize an MPI environment in jupyter lab. Similarly, in the following code blocks the lines containing %%px can be ignored. To run your own script in an MPI environment simply execute it with mpirun -n <n_tasks> python3 <script_name>.

# IGNORE THIS BLOCK
import ipyparallel
cluster = ipyparallel.Client()

Initialize MPI varibales, cartesian tensor operators and external field:

%%px

from mpi4py import MPI
from richmol.field import CarTens
from richmol.convert_units import Debye_x_Vm_to_invcm

# INITIALIZE MPI VARIABLES
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
world_size = comm.Get_size()
if rank == 0:
    print('world size : ', world_size)

# INITIALIZE CARTESIAN TENSOR OPERATORS
richmol_file = "h2o_p48_j40_emax4000_rovib.h5"
h0 = CarTens(richmol_file, name='h0')
mu = CarTens(richmol_file, name='dipole') \
    * (-0.5 * Debye_x_Vm_to_invcm())

# INITIALIZE EXTERNAL DC FIELD
field = lambda t : [0, 0, t * 1e5 ] # (V/m)

[stdout:0] world size :  16

A TDSE solver and initial state vectors, spanning the space of field-free states, are initialized. This time, however, initial state vectors are distributed over the MPI environment:

%%px

from richmol.tdse import TDSE

# INITIALIZE TDSE AND STATES
tdse = TDSE(t_start=0, t_end=1, dt=0.01, t_units="ps", enr_units="invcm")
vecs = tdse.init_state(h0, temp=None)

# DISTRIBUTE INITIAL STATES
n_vecs = int(len(vecs) / world_size + 1)
if not rank == (world_size - 1):
    vecs = vecs[rank * n_vecs : (rank + 1) * n_vecs]
else:
    vecs = vecs[rank * n_vecs :]
print("process with `rank` = '{}' propagates {} states".format(rank, len(vecs)))

[stdout:0] process with `rank` = '0' propagates 563 states
[stdout:1] process with `rank` = '1' propagates 563 states
[stdout:2] process with `rank` = '2' propagates 563 states
[stdout:3] process with `rank` = '3' propagates 563 states
[stdout:4] process with `rank` = '4' propagates 563 states
[stdout:5] process with `rank` = '5' propagates 563 states
[stdout:6] process with `rank` = '6' propagates 563 states
[stdout:7] process with `rank` = '7' propagates 563 states
[stdout:8] process with `rank` = '8' propagates 563 states
[stdout:9] process with `rank` = '9' propagates 563 states
[stdout:10] process with `rank` = '10' propagates 563 states
[stdout:11] process with `rank` = '11' propagates 563 states
[stdout:12] process with `rank` = '12' propagates 563 states
[stdout:13] process with `rank` = '13' propagates 563 states
[stdout:14] process with `rank` = '14' propagates 563 states
[stdout:15] process with `rank` = '15' propagates 550 states

The intial states are propagated in time:

%%px

import time

comm.Barrier()
time_0 = time.time()
for i, t in enumerate(tdse.time_grid()):
    vecs_, _ = tdse.update(
        mu.field(field(t)), H0=h0, vecs=vecs, matvec_lib='scipy'
    )
    vecs = vecs_
comm.Barrier()
runtime = time.time() - time_0
if rank == 0:
    print('runtime with MPI : {} sec'.format(round(runtime, 3)))

[stdout:0] runtime with MPI : 22.315 sec