<jupyter_start><jupyter_text>Non-adiabatic dynamics This tutorial shows how to run non-adiabatic dynamics with a trained model using [Tully's surface hopping](https://doi.org/10.1063/1.459170).First let's import dependencies:<jupyter_code>import os
import json
from nff.md.tully.io import get_results, load_json, get_atoms
from nff.md.tully.dynamics import NeuralTully, CombinedNeuralTully
from ase.io.trajectory import Trajectory
import nglview as nv<jupyter_output><empty_output><jupyter_text>Now we'll find a trained model. The trained azobenzene models can be found in `NeuralForceField/models/azo_derivatives`. The sub-folders are for diabatic and adiabatic models, trained either with the full set of geometries, or with 40 species held out. There are also three models trained with different splits and different initialization from random seeds:<jupyter_code>print(os.listdir('../models/azo_derivatives'))
print(os.listdir('../models/azo_derivatives/all_diabatic'))<jupyter_output>['all_diabatic', 'holdout_diabatic', 'holdout_adiabatic']
['seed_0', 'seed_2', 'seed_1']<jupyter_text>We'll use the diabatic model trained on all species, with seed 0: `../models/azo_derivatives/all_diabatic/seed_0`.  TullyThe script for Tully's surface hopping is `NeuralForceField/nff/md/tully/dynamics.py`. If you run the script and supply the path of a JSON parameter file, it will do the rest for you. Here we'll go through some parameters to give, and show a little of what goes on behind the scenes.We'll have to define `ground_params`, `tully_params`, and `all_params`. The first is for parameters in the ground state MD simulation, and the second for Tully's surface hopping. The third is for some remaining parameters, which we'll explain below.Let's define `ground_params`:<jupyter_code>ground_params = {'ttime': 50, # tau = ttime * dt is the relaxation time
                 'logfile': 'ground.log', # log file for ground state MD
                 'max_time': 200, # total time in fs
                 'savefile': 'ground.trj', # output file with saved geometries
                 'timestep': 0.5, # dt in fs
                 'equil_time': 100, # ignore this amount of time (fs) when sampling
                                  # geoms for NAMD
                 'thermostat': 'nosehoover', # use the Nose Hoover thermostat
                 'loginterval': 5, # log the energy and save geoms every 5 steps
                 'temperature': 300, # temperature in Kelvin
                 'maxwell_temp': 300, # initialization temperature from
                                      # Maxwell-Boltzmann distribution
                 'nbr_update_period': 10, # update the neighbor list every 10 steps
                 'cutoff': 5.0, # neighbor list cutoff in Angstrom 
                 'cutoff_skin': 2.0, # extra distance added to cutoff when updating
                                     # neighbor list, to account for atoms coming into
                                     # the 5 A sphere between updates 
                 'model_type': 'PainnDiabat', # model class we're using
                 'device': 1, # GPU to use. Set to "cpu" if you don't have a GPU
                }<jupyter_output><empty_output><jupyter_text>Now let's do `tully_params`:<jupyter_code>diabat_keys = [['d_00', 'd_01', 'd_02'],
               ['d_01', 'd_11', 'd_12'],
               ['d_02', 'd_12', 'd_22']]

tully_params = {
                 'dt': 0.5, # dt in fs
                 'cutoff': 5.0, # neighbor list cutoff in Angstrom 
                 'hop_eqn': 'sharc', # use the SHARC equation for computing the 
                                     # hopping probability
                 'num_trj': 10, # number of surface hopping trajectories
                 'max_time': 200, # maximum time in fs
                 'batch_size': 5, # number of trajectories to batch together 
                                  # so that inference is done in parallel
                 'num_states': 2, # number of adiabatic states
                 'cutoff_skin': 2.0, # extra distance added to cutoff when updating
                                     # neighbor list, to account for atoms coming into
                                     # the 5 A sphere between updates 
                 'decoherence': {'name': 'truhlar'}, # use Truhlar's decoherence method
                 'max_gap_hop': None, # maximum gap, above which hops cannot occur
                 'save_period': 10, # save every 10 steps (5 fs)
                 'initial_surf': 1, # start on the first excited adiabatic state
                 'initial_time': 0.0, # start at t = 0
                 'elec_substeps': 25, # 25 electronic steps for every nuclear step
                 'reload_ground': False, # If the ground state trajectory is already
                                        # complete from a previous run, then do not load
                                        # it instead of re-running
                 'diabat_propagate': True, # propagate in the diabatic basis
                 'simple_vel_scale': True, # after a hop, simply rescale the velocity
                                           # to conserve energy, instead of scaling in
                                           # NACV direction
                 'nbr_update_period': 10, # update neighbor list every 10 steps
                 'explicit_diabat_prop': True, # propagate in diabatic basis,                                 
                 'device': 1, # GPU to use. Set to "cpu" if you don't have a GPU
                 'diabat_keys': diabat_keys # names of diabatic energies predicted
                                # by model
                
    
                
}<jupyter_output><empty_output><jupyter_text>Lastly we'll define `all_params`, which has the starting coordinates and the model path:<jupyter_code>with open('data/azo_coords.json', 'r') as f:
    coords = json.load(f)

all_params = {"coords": coords, # starting geometry of the molecule
              'model_path': '../models/azo_derivatives/all_diabatic/seed_0'
             }<jupyter_output><empty_output><jupyter_text>When we run the script from the command line, it parses these three dictionaries from a file and makes an instance of `CombinedNeuralTully`, like this:<jupyter_code>atomsbatch = get_atoms(ground_params=ground_params,
                      all_params=all_params)

tully_params.update({"model_path": all_params['model_path']})

neural_tully = CombinedNeuralTully(atoms=atomsbatch,
                                   ground_params=ground_params,
                                   tully_params=tully_params)<jupyter_output>/home/saxelrod/miniconda3/envs/nff/lib/python3.7/site-packages/ase/md/md.py:48: FutureWarning: Specify the temperature in K using the 'temperature_K' argument
  warnings.warn(FutureWarning(w))<jupyter_text>The first thing it does is make an `AtomsBatch` object. This is a class we made that inherits from the ASE `Atoms` object, and has othe functionalities to interface with NFF. Then an instance of `CombinedNeuralTully` is made. This object runs ground-state dynamics to generate starting frames and velocities, and then performs surface hopping. It's called "combined" because it combines ground- and excited-state dynamics.For an example of how you would use this script in practice, check out `data/tully_info.json`. If you run```bashconda activate nffpython ../nff/md/tully/dynamics.py --params_file data/tully_info.json```then you should be able to peform neural Tully in one line. Note that in `tully_info.json`, the `all_params` part of the dictionary is its body, i.e. everything that doesn't have the key `ground_params` or `tully_params`.**Note**: it is very important to have PyTorch version >=1.9 installed. PyTorch is used for exponentiation of complex matrices, which the older versions could not do. Instead of raising an error, the old versions silently return garbage!<jupyter_code># This takes a few minutes to run
neural_tully.run()<jupyter_output>../nff/nn/modules/diabat.py:141: UserWarning: torch.symeig is deprecated in favor of torch.linalg.eigh and will be removed in a future PyTorch release.
The default behavior has changed from using the upper triangular portion of the matrix by default to using the lower triangular portion.
L, _ = torch.symeig(A, upper=upper)
should be replaced with
L = torch.linalg.eigvalsh(A, UPLO='U' if upper else 'L')
and
L, V = torch.symeig(A, eigenvectors=True)
should be replaced with
L, V = torch.linalg.eigh(A, UPLO='U' if upper else 'L') (Triggered internally at  /opt/conda/conda-bld/pytorch_1623448224956/work/aten/src/ATen/native/BatchLinearAlgebra.cpp:2500.)
  ad_energies, u = torch.symeig(d_mat, True)<jupyter_text>We can view the ground-state log file:<jupyter_code>with open('ground.log', 'r') as f:
    ground_log = f.read()
print(ground_log)<jupyter_output>Time[ps]      Etot[eV]     Epot[eV]     Ekin[eV]    T[K]
0.0000           0.1020      -0.8128       0.9147   208.1
0.0025           0.1067      -0.5076       0.6143   139.8
0.0050           0.1085      -0.3109       0.4194    95.4
0.0075           0.1146      -0.2349       0.3495    79.5
0.0100           0.1211      -0.4860       0.6071   138.2
0.0125           0.1381      -0.3638       0.5020   114.2
0.0150           0.1522      -0.3410       0.4932   112.2
0.0175           0.1703      -0.2687       0.4390    99.9
0.0200           0.1911      -0.4260       0.6171   140.4
0.0225           0.2206      -0.2439       0.4645   105.7
0.0250           0.2459      -0.2879       0.5339   121.5
0.0275           0.2795      -0.2799       0.5594   127.3
0.0300           0.3203      -0.3854       0.7057   160.6
0.0325           0.3676      -0.2013       0.5689   129.5
0.0350           0.4128      -0.2612       0.6741   153.4
0.0375           0.4683      -0.2329       0.7012   159.6
0.0400         [...]<jupyter_text>We see that all energies fluctuate, as kinetic energy is being added into the system fo the thermostat. The temperature also varies, and over enough time it will average out to 300 K. To get the actual geometries, energies, and forces, we can load the trajectory file. And we can visualize it with `nglview`:<jupyter_code>trj = Trajectory('ground.trj')
nv.show_asetraj(trj)<jupyter_output><empty_output><jupyter_text>We can also view the excited state log:<jupyter_code>with open('tully.log', 'r') as f:
    tully_log = f.read()
print(tully_log)<jupyter_output>Time [fs]         State 0         State 1             |c|       Hop prob. 
0.5                 0.0000%       100.0000%         1.0000         0.0000
1.0                 0.0000%       100.0000%         1.0000         0.0000
1.5                 0.0000%       100.0000%         1.0000         0.0000
2.0                 0.0000%       100.0000%         1.0000         0.0000
2.5                 0.0000%       100.0000%         1.0000         0.0000
3.0                 0.0000%       100.0000%         1.0000         0.0000
3.5                 0.0000%       100.0000%         1.0000         0.0000
4.0                 0.0000%       100.0000%         1.0000         0.0000
4.5                 0.0000%       100.0000%         1.0000         0.0000
5.0                 0.0000%       100.0000%         1.0000         0.0000
5.5                 0.0000%       100.0000%         1.0000         0.0000
6.0                 0.0000%       100.0000%         1.0000         0.0000
6.5                 0.0000%       100[...]<jupyter_text>This log shows us the population in the different states at each time, the norm of the electronic coefficient vector $c$, and the maximum hopping probability of any of the trajectories. We see that all of the trajectories have returned to the ground state by 200 fs.To get the geometries, forces, etc., we can use the `NeuralTully` class method `from_pickle`:<jupyter_code>state_dicts, trjs = NeuralTully.from_pickle('tully.pickle')<jupyter_output><empty_output><jupyter_text>The output contains `state_dicts`, a set of dictionaries with information about each trajectory for each time step, and `trjs`, a set of ASE trajectories. We can visualize one of the trajectories:<jupyter_code>trj = trjs[1]
nv.show_asetraj(trj)<jupyter_output><empty_output><jupyter_text>We can also take a look at `state_dicts`:<jupyter_code>print(len(state_dicts))
print(list(state_dicts[0].keys()))<jupyter_output>49
['nxyz', 'nacv', 'force_nacv', 'energy', 'forces', 'H_d', 'U', 't', 'vel', 'c', 'T', 'surfs']<jupyter_text>We see that there is one dictionary for each saved time step (all time steps in which a hop occurs are automatically saved, even if you only asked the geoms to be saved every X steps). Each dictionary has various properties, each of shape `num_trj  x ...`, so that the first dimension divides properties by trajectory. For example:<jupyter_code>state_dict = state_dicts[0]
print('nxyz: ', state_dict['nxyz'].shape) # coordinates 
print('energy: ', state_dict['energy'].shape) # energy
print('forces: ', state_dict['forces'].shape) # forces on each surface
print('H_d: ', state_dict['H_d'].shape) # diabatic Hamiltonian
print('U: ', state_dict['U'].shape) # unitary transformation from diabatic to adiabatic 
print('vel: ', state_dict['vel'].shape) # velocities
print('c: ', state_dict['c'].shape) # c electronic wavefunction vector 
print('surfs: ', state_dict['surfs'].shape) # current surface


print(state_dict['force_nacv']) # not computed because we did diabatic propagation
print(state_dict['T']) # T matrix (NACV off-diagonal components of Hamiltonian)
                       # also not computed because we did diabatic propagation<jupyter_output>nxyz:  (10, 34, 4)
energy:  (10, 2)
forces:  (10, 2, 34, 3)
H_d:  (10, 3, 3)
U:  (10, 3, 3)
vel:  (10, 34, 3)
c:  (10, 3)
surfs:  (10,)
None
None
{'fraction_non_alphanumeric': 0.09326349483486104, 'fraction_numerical': 0.06219991270187691, 'mean_word_length': 2.3415167719980556, 'pattern_counts': {'":': 2, '<': 50, '<?xml version=': 0, '>': 51, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '14', 'dir': 'jupyter-structured-clean-dedup', 'id': '475506', 'license': 'MIT', 'n_tokens_mistral': 4684, 'n_tokens_neox': 4039, 'n_words': 1376, 'path': 'tutorials/10_tully_namd.ipynb', 'repo_name': 'jkaraguesian/NeuralForceField', 'size': '13746'}