Algorithm Development

In this guide, we will discuss how to make in-place modifications to Algorithms. Algorithm development is a more advanced topic and requires a good understanding of the underlying steps of the algorithm so we will not go into detail (in this guide) about how to develop a new algorithm from scratch. Instead, we will focus on making changes to existing algorithms.

Before we proceed, let’s discuss the structure of an algorithm in QC Lab. An algorithm in QC Lab is a Python class that inherits from the Algorithm class in the qc_lab.algorithm module. The built-in algorithms are found in the qc_lab.algorithms module and presently contain qc_lab.algorithms.MeanField and qc_lab.algorithms.FewestSwitchesSurfaceHopping.

We can start by importing the MeanField algorithm from the qc_lab.algorithms module:

from qc_lab.algorithms import MeanField

Each algorithm consists of three lists of functions which are referred to as “recipes”, the functions themselves are referred to as “tasks”.

# the initialization recipe
print(MeanField.initialization_recipe)

# the update recipe
print(MeanField.update_recipe)

# the output recipe
print(MeanField.output_recipe)

As the name implies, the initialization_recipe initializes all the variables required for the algorithm. The update_recipe updates the variables at each time step, and the output_recipe is used to output the results of the algorithm.

In addition to the recipes, a list of variable names is needed to specify which variables the algorithm will store in the Data object.

# the variables that the algorithm will store
print(MeanField.output_variables)

Adding output obvservables

To add an additional variable to the output of an algorithm we must define a task that calculates the variable and add it to the output recipe.

Linear response functions

For example, let’s calculate a linear response function that can be used to calculate the absorption spectrum of a system. Mathematically we will calculate

\[R(t) = \langle \psi(0) \vert \psi(t)\rangle\]

where \(\vert \psi(t)\rangle\) is the diabatic wavefunction at time \(t\).

def update_response_function(sim, parameters, state, **kwargs):
    # First get the diabatic wavefunction.
    wf_db = state.wf_db
    # If we are at the first timestep we can store the diabatic wavefunction in the parameters object
    if sim.t_ind == 0:
        parameters.wf_db_0 = np.copy(wf_db)
    # Next calculate the response function and store it in the state object.
    state.response_function = np.sum(np.conj(parameters.wf_db_0) * wf_db, axis=-1)
    return parameters, state

Next we can add this task to the output recipe.

MeanField.output_recipe.append(update_response_function)

Finally we can add the relevant variable name to the output_variables list.

MeanField.output_variables.append('response_function')

We can then run a simulation and calculate the corresponding spectral function,

from qc_lab import Simulation
from qc_lab.dynamics import parallel_driver_multiprocessing
from qc_lab.models import SpinBoson

# instantiate a simulation
sim = Simulation()
print('default simulation settings: ', sim.default_settings)

# change settings to customize simulation
sim.settings.num_trajs = 1000
sim.settings.batch_size = 250
sim.settings.tmax = 50
sim.settings.dt = 0.01

# instantiate a model
sim.model = SpinBoson({'l_reorg': 0.2})
print('default model constants: ', sim.model.default_constants) # print out default constants

# instantiate an algorithm
sim.algorithm = MeanField()
print('default algorithm settings: ', sim.algorithm.default_settings) # print out default settings



# define an initial diabatic wavefunction
sim.state.wf_db = np.array([1, 0], dtype=complex)

# run the simulation
data = parallel_driver_multiprocessing(sim, num_tasks=4)

# plot the data.
print('calculated quantities:', data.data_dic.keys())
response_function = data.data_dict['response_function']
time = sim.settings.tdat_output
plt.plot(time, np.real(response_function), label='R(t)')
plt.xlabel('time')
plt.ylabel('response function')
plt.legend()
plt.show()

plt.plot(np.real(np.roll(np.fft.fft(response_function), len(time)//2)))
plt.xlabel('freq')
plt.ylabel('absorbance')
plt.show()

Adiabatic populations

Next, let’s calculate the adiabatic populations of the system as is sometimes done in scattering problems. Obviously these populations will only have a well-defined meaning in regimes with no nonadiabatic coupling.

def update_adiabatic_populations(sim, parameters, state, **kwargs):
    # First get the Hamiltonian and calculate its eigenvalues and eigenvectors.
    H = state.h_quantum # this is the quantum plus quantum-classical Hamiltonian.
    # Next obtain its eigenvalues and eigenvectors.
    evals, evecs = np.linalg.eigh(H)
    # Now calculate the adiabatic wavefunction.
    wf_adb = np.einsum('tji,tj->ti', np.conj(evecs), state.wf_db)
    # Finally calculate the populations (note that we do not sum over the batch, this is done internally by QC Lab).
    pops_adb = np.abs(wf_adb)**2
    # Store the populations in the state object.
    state.pops_adb = pops_adb
    return parameters, state

Next we can add this task to the output recipe.

MeanField.output_recipe.append(update_adiabatic_populations)

Finally we can add the relevant variable name to the output_variables list.

MeanField.output_variables.append('pops_adb')

We can then run a simulation and plot the populations. Note that since the spin-boson model is always in a coupling regime these populations will not have a well-defined meaning.

from qc_lab import Simulation
from qc_lab.dynamics import serial_driver
from qc_lab.models import SpinBoson

# Instantiate a simulation.
sim = Simulation()
print('default simulation settings: ', sim.default_settings)

# Change settings to customize simulation.
sim.settings.num_trajs = 100
sim.settings.batch_size = 100
sim.settings.tmax = 25
sim.settings.dt = 0.01

# Instantiate a model.
sim.model = SpinBoson()
print('default model constants: ', sim.model.default_constants) # print out default constants

# Instantiate an algorithm.
sim.algorithm = MeanField()
print('default algorithm settings: ', sim.algorithm.default_settings) # print out default settings



# Define an initial diabatic wavefunction.
sim.state.wf_db = np.array([1, 0], dtype=complex)

# Run the simulation.
data = serial_driver(sim)

# Plot the data.
print('calculated quantities:', data.data_dic.keys())
classical_energy = data.data_dict['classical_energy']
quantum_energy = data.data_dict['quantum_energy']
populations = np.real(np.einsum('tii->ti', data.data_dict['dm_db']))
adiabatic_populations = np.real(data.data_dict['pops_adb'])
time = sim.settings.tdat_output
plt.plot(time, adiabatic_populations[:,0], label='adiabatic state 0')
plt.plot(time, adiabatic_populations[:,1], label='adiabatic state 1')
plt.xlabel('time')
plt.ylabel('population')
plt.legend()
plt.show()

Note

In the above code we chose to modify the MeanField class itself rather than an instance of it. This can lead to troublesome behavior in a Jupyter notebook where the class will not be reloaded if the cell is rerun. Restarting the kernel will fix this issue. Otherwise one can modify an instance of the class by creating a new instance and modifying it.

Modifying algorithm behavior

In the same way that we could modify the output recipe, it is possible to modify the initialization and update recipes in the same way. We will not go into detail on how to do this here but the process is the same as for the output recipe (except there is no output variable in those cases).