Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import torchani
- import torch
- from simtk import unit
- import numpy as np
- import matplotlib.pyplot as plt
- hartree_to_kJ_mol = 2625.499638
- nm_to_angstroms = (1.0 * unit.nanometer) / (1.0 * unit.angstrom)
- device = torch.device('cpu')
- model = torchani.models.ANI1ccx()
- species = model.species_to_tensor('CHHHH').to(device).unsqueeze(0)
- def ANI1ccx_energy(x):
- """
- Parameters
- ----------
- x : numpy array, unit'd
- coordinates
- Returns
- -------
- energy : float, unit'd
- energy, with unit of kJ/mol attached
- """
- assert(type(x) == unit.Quantity)
- x_in_nm = x.value_in_unit(unit.nanometer)
- coordinates = torch.tensor([x_in_nm],
- requires_grad=False, device=device, dtype=torch.float32)
- # convert from nm to angstroms
- coordinates_in_angstroms = coordinates * nm_to_angstroms
- _, energy_in_hartree = model((species, coordinates_in_angstroms))
- # convert energy from hartrees to kJ/mol
- energy_in_kJ_mol = energy_in_hartree * hartree_to_kJ_mol
- return float(energy_in_kJ_mol.detach()) * unit.kilojoule_per_mole
- def ANI1ccx_energy_ensemble(x):
- """
- Parameters
- ----------
- x : numpy array, unit'd
- coordinates
- Returns
- -------
- energies : numpy array, unit'd
- 8 energies from the ensemble, with units of kJ/mol attached
- """
- assert(type(x) == unit.Quantity)
- x_in_nm = x.value_in_unit(unit.nanometer)
- coordinates = torch.tensor([x_in_nm],
- requires_grad=False, device=device, dtype=torch.float32)
- # convert from nm to angstroms
- coordinates_in_angstroms = coordinates * nm_to_angstroms
- energies_in_kJ_mol = [float(hartree_to_kJ_mol * m((species, coordinates_in_angstroms))[1].detach()) for m in model]
- return np.array(energies_in_kJ_mol) * unit.kilojoule_per_mole
- scale_factors = np.linspace(0,2,1000)
- tetrahedral_geometry = np.array([
- (0, 0, 0),
- (1, 1, 1),
- (-1, -1, 1),
- (-1, 1, -1),
- (1, -1, -1),
- ], dtype=float)
- tetrahedral_geometry /= np.linalg.norm(tetrahedral_geometry[0] - tetrahedral_geometry[1])
- # lowest energy found
- bent_optimum = np.array([
- [ 0. , 0. , 0. ],
- [ 0.64702535, 0.64647233, 0.8811238 ],
- [-0.64734113, -0.6475463 , 0.8808765 ],
- [-0.7168874 , 0.7172707 , -0.33484012],
- [ 0.7182246 , -0.71852016, -0.33497408]
- ]) # in angstroms
- lowest_energy = ANI1ccx_energy(bent_optimum * unit.angstroms)
- # scaled so that the average bond length is 1 angstrom
- normalized_bent_optimum = bent_optimum / np.mean(np.linalg.norm(bent_optimum[1:], axis=1))
- # energy scan from global optimum (bent geometry), with ensemble uncertainty
- bent_scan = np.array([normalized_bent_optimum * scale_factor for scale_factor in scale_factors])
- bent_energies = [ANI1ccx_energy_ensemble(g * unit.angstrom) for g in bent_scan]
- bent_energy_array = np.array([(e - lowest_energy) / unit.kilojoule_per_mole for e in bent_energies])
- bent_means = np.mean(bent_energy_array, 1)
- bent_stdevs = np.std(bent_energy_array, 1)
- # energy scan from tetrahedral geometry, with ensemble uncertainty
- tetrahedral_scan = np.array([tetrahedral_geometry * scale_factor for scale_factor in scale_factors])
- tetra_energies = [ANI1ccx_energy_ensemble(g * unit.angstrom) for g in tetrahedral_scan]
- tetra_energy_array = np.array([(e - lowest_energy) / unit.kilojoule_per_mole for e in tetra_energies])
- tetra_means = np.mean(tetra_energy_array, 1)
- tetra_stdevs = np.std(tetra_energy_array, 1)
- def average_bond_length_in_snapshot(xyz):
- return np.mean(np.linalg.norm(xyz[1:], axis=1))
- # plot
- plt.figure()
- ax = plt.subplot(1,1,1)
- ax.spines['top'].set_visible(False)
- ax.spines['right'].set_visible(False)
- plt.ylim(-20,400)
- plt.xlim(0.85,1.35)
- plt.xlabel('C-H bond length (Å)')
- plt.ylabel('$\Delta E \pm$ stddev (kJ/mol)\nrelative to global minimum')
- plt.title('methane energy scans: varying C-H bond length')
- # tetrahedral
- # x: average C-H bond length
- tetrahedral_x = [average_bond_length_in_snapshot(xyz) for xyz in tetrahedral_scan]
- # y: mean energy +/- standard deviation
- plt.plot(tetrahedral_x, tetra_means, label='scan from tetrahedron')
- plt.fill_between(tetrahedral_x, tetra_means - tetra_stdevs, tetra_means + tetra_stdevs, alpha=0.5)
- # bent
- # x: average C-H bond length
- bent_x = [average_bond_length_in_snapshot(xyz) for xyz in bent_scan]
- # y: mean energy +/- standard deviation
- plt.plot(bent_x, bent_means, label='scan from global minimum (bent)')
- plt.fill_between(bent_x, bent_means - bent_stdevs, bent_means + bent_stdevs, alpha=0.5)
- # horizontal line indicating global minimum
- plt.hlines(0, 0,2,
- linestyles='dotted', linewidth=1, label='global minimum (bent)')
- # clip x and y limits
- plt.ylim(-150,400)
- plt.xlim(0.85,1.35)
- plt.legend(loc='lower left')
- plt.tight_layout()
- plt.savefig('methane_energy_scans_with_uncertainties.png', dpi=300, bbox_inches='tight')
- plt.close()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement