from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem.Draw import IPythonConsole
import py3Dmol
import torch
import torchani
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
from psikit import Psikit
pk = Psikit()
smiles = 'c1ccccc1-c2c(C)cncc2'
pk.read_from_smiles(smiles)
pk.rdkit_optimize()
%time pk.energy(basis_sets='WB97X/6-31g*')
CPU times: user 1min 10s, sys: 1.79 s, total: 1min 12s Wall time: 18.5 s
-518.5098949584047
IPythonConsole.drawMol3D(pk.mol)
You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol
model = torchani.models.ANI2x(periodic_table_index=True).to(device)
def mol2arr(mols, device=device):
coordinates = []
spices = []
for mol in mols:
pos = mol.GetConformer().GetPositions().tolist()
atomnums = [a.GetAtomicNum() for a in mol.GetAtoms()]
coordinates.append(pos)
spices.append(atomnums)
coordinates = torch.tensor(coordinates,
requires_grad=True,
device=device)
species = torch.tensor(spices, device=device)
return coordinates, species
coordinates, species = mol2arr([pk.mol], device)
%time energy = model((species, coordinates)).energies
derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
force = -derivative
CPU times: user 16.5 ms, sys: 15 µs, total: 16.5 ms Wall time: 16.2 ms
print('Energy:', energy.item())
print('Force:', force.squeeze())
Energy: -518.5092850781618 Force: tensor([[ 1.3592e-02, 3.5025e-03, 2.0906e-02], [-1.4324e-02, -1.6894e-04, 1.4033e-02], [-1.3329e-02, -3.7380e-03, 2.3463e-04], [-9.2622e-03, -5.9881e-03, -1.5998e-02], [ 2.1978e-02, 4.3879e-03, -1.6717e-02], [ 1.4468e-02, 7.4550e-03, 3.5641e-03], [-2.8785e-02, -9.5195e-03, -5.3672e-03], [-2.6472e-03, -3.8259e-03, 8.0281e-04], [-9.3185e-04, 1.7754e-02, -1.3840e-03], [ 6.6303e-02, -7.4541e-02, -4.3813e-02], [-1.2796e-02, -1.2576e-03, 1.3008e-03], [-2.4707e-03, 7.7882e-02, 3.8845e-02], [-1.5166e-02, 6.1120e-03, 1.5764e-03], [-5.9885e-03, -2.7664e-03, -4.3181e-03], [ 2.6414e-03, -5.7588e-06, -5.2600e-03], [ 5.9540e-03, 1.8318e-03, 4.1510e-04], [ 1.0910e-03, 1.0532e-03, 4.8398e-03], [-7.2534e-03, -1.1427e-03, 3.0659e-03], [-1.5687e-03, 1.0146e-02, -1.6901e-02], [-1.1689e-02, -1.0652e-02, 1.5877e-02], [ 1.5274e-02, -1.3123e-02, 6.3923e-03], [-9.5655e-03, -1.6000e-03, -5.9007e-04], [-1.3352e-02, -2.3642e-03, -2.1185e-04], [ 7.8282e-03, 5.6739e-04, -1.2904e-03]], device='cuda:0')
force
tensor([[[ 1.3592e-02, 3.5025e-03, 2.0906e-02], [-1.4324e-02, -1.6894e-04, 1.4033e-02], [-1.3329e-02, -3.7380e-03, 2.3463e-04], [-9.2622e-03, -5.9881e-03, -1.5998e-02], [ 2.1978e-02, 4.3879e-03, -1.6717e-02], [ 1.4468e-02, 7.4550e-03, 3.5641e-03], [-2.8785e-02, -9.5195e-03, -5.3672e-03], [-2.6472e-03, -3.8259e-03, 8.0281e-04], [-9.3185e-04, 1.7754e-02, -1.3840e-03], [ 6.6303e-02, -7.4541e-02, -4.3813e-02], [-1.2796e-02, -1.2576e-03, 1.3008e-03], [-2.4707e-03, 7.7882e-02, 3.8845e-02], [-1.5166e-02, 6.1120e-03, 1.5764e-03], [-5.9885e-03, -2.7664e-03, -4.3181e-03], [ 2.6414e-03, -5.7588e-06, -5.2600e-03], [ 5.9540e-03, 1.8318e-03, 4.1510e-04], [ 1.0910e-03, 1.0532e-03, 4.8398e-03], [-7.2534e-03, -1.1427e-03, 3.0659e-03], [-1.5687e-03, 1.0146e-02, -1.6901e-02], [-1.1689e-02, -1.0652e-02, 1.5877e-02], [ 1.5274e-02, -1.3123e-02, 6.3923e-03], [-9.5655e-03, -1.6000e-03, -5.9007e-04], [-1.3352e-02, -2.3642e-03, -2.1185e-04], [ 7.8282e-03, 5.6739e-04, -1.2904e-03]]], device='cuda:0')