yanzewu / pywp

Quantum wavepacket simulation using split operator method

PyWP

Efficient and customizable code to run Cartesian wavepacket dynamics using the split operator method. Supports arbitrary number of electronic states and nuclear degrees of freedom. Supports GPU.

Requires python >= 3.7, numpy >= 1.9 and (for GPU only) cupy >= 9.5.

Tutorial

PyWP handles these parameters in two steps: First, preprocess() generate necessary caches and then propagate() does the actual propagation. A runnable script is found at test/wavepacket.py.

Here is an example of running Tully1:

(1) Define the necessary ingredients. To run a wavepacket simulation, we need (a) A "canvas" or grid; (b) An initial condition; (c) A Hamiltonian;

import pywp

# Defining the grid
grid = pywp.mgrid[-10:10:0.01]

# Defining the wavepacket (normalization is not needed)
wavepacket = lambda R: np.exp(-(R[0]+5)**2 + 1j*P0*R[0])

# Defining the potential
tully1 = lambda R: pywp.build_pe_tensor(
    0.01*(1 - np.exp(-1.6*R[0])) * (R[0] >= 0) 
        - 0.01*(1 - np.exp(1.6*R[0])) * (R[0] < 0), # V00
    0.005*np.exp(-R[0]**2), # V01
    symmetric=True  # copy V11 = -V00
 )

(2) Do the preprocessing:

para = pywp.preprocess(
    tully1, grid, wavepacket, 
    c0=[1,0],   # initial amplitude
    M=2000,     # mass
    dt=0.1,     # timestep
)

(3) Define the collectors we want to see. Here we want to collected the transmitted population on individual electronic states and save the snapshots.

pop_trans = []

define pop_trans_collector(para:pywp.PhysicalParameter, checkpoint:pywp.CheckPoint):
    populations.append(
        np.sum(np.abs(checkpoint.psiR)**2 * (para.R[0] > 0)[...,None], axis=0) * para.dA
    )

snapshot_writer = pywp.snapshot.SnapshotWriter('tully1.snapshot')

(4) Do the propagation:

pywp.propagate(para, 
                nstep=12000,        # Total step for propagation
                output_step=1000,   # Step interval per output
                on_output=[pop_trans_collector, snapshot_writer],   # Output functions
                checkend=True,      # Checkend
                boundary=lambda R: R[0] < 9,    # Checkend boundaries
            )

snapshot_writer.close() # Don't forget to do this

print(pop_trans)    # should be a list of array of 2 elements

(5) To view your running result:

snapshot = pywp.snapshot.load_file('tully1.snapshot')
pywp.visualize.visualize_snapshot_1d(snapshot)  # press '.' to forward, ',' to backward and 'q' to quit

Snapshots

The wavefunction can be retrieved by snapshot.get_snapshot(). The position and momentum grids can be retrieved by snapshot.get_R_grid() and snapshot.get_P_grid().

Visualization

PyWP comes with a powerful visualizing support, which make it easy to test custom potentials. Most functions are in pywp.visualize.

• visualize_pe_1d(), visualize_pe_2d() and visualize_pe_2d_surf() plot the potential surface in one or two specific dimension(s).
• visualize_snapshot_1d(), visualize_snapshot_2d() plot the wavepacket in one or two specific dimension(s).
• imshow(), imshow_multiple(), surf_multiple() are wrappers around matplotlib functions, which can conveniently plot 2D data.

python -m pywp.tools.show1d and python -m pywp.tools.show2d are command line wrappers for visualizing snapshots.

Scattering Calculation (Experimental)

An experimental implementation of 1D scattering calculation is in pywp.scattering.scatter1d(). The usage is almost straightforward.