From 8a22efd2b781148104fa8cd9cfcc9a7b03558809 Mon Sep 17 00:00:00 2001 From: Alexander Humeniuk Date: Tue, 5 Jan 2021 16:37:48 +0100 Subject: [PATCH 1/9] The new predictor `torchtools_hessian.GDMLPredict(...)` can evaluate second order derivatives of the potential energy (Hessians) in addition to energies and gradients. `test_torchtools_hessian.py` demonstrates that the new implementation of GDMLPredict(...) gives the same energies and gradients as the original GDMLTorchPredict(...) . The analytical Hessians are compared to numerical Hessian computed with ASE. Certain semiclassical propagators such as the Herman-Kluk propagator require a local harmonic approximation around each trajectory. Computing the Hessian numerically for each time step would be too inaccurate and time-consuming. The analytical Hessian calculation is roughly 5-10 times more expensive than a gradient calculation. --- sgdml/test_torchtools_hessian.py | 146 +++++++++++++++++++ sgdml/torchtools_hessian.py | 239 +++++++++++++++++++++++++++++++ sgdml/train.py | 2 +- 3 files changed, 386 insertions(+), 1 deletion(-) create mode 100755 sgdml/test_torchtools_hessian.py create mode 100644 sgdml/torchtools_hessian.py diff --git a/sgdml/test_torchtools_hessian.py b/sgdml/test_torchtools_hessian.py new file mode 100755 index 0000000..6dc907b --- /dev/null +++ b/sgdml/test_torchtools_hessian.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python +# coding: utf-8 +""" +test torch implementation of machine-learned sGDML model (energies, gradients and Hessians) +""" +model_file = 'models/ethanol.npz' +geometry_file = 'geometries/ethanol.xyz' + +import numpy as np +import numpy.linalg as la +import torch +import logging +import time + +from sgdml.utils import io + +from sgdml.torchtools_hessian import GDMLPredict +from sgdml.torchtools import GDMLTorchPredict + +# # Logging +logger = logging.getLogger(__name__) + +# GPU or CUDA? +torch.set_default_dtype(torch.float64) +if torch.cuda.is_available(): + logger.info("CUDA available") + device = torch.device("cuda") +else: + device = torch.device('cpu') + +# load model fitted to ground state forces +model = np.load(model_file, allow_pickle=True) +# reference implementation +gdml_ref = GDMLTorchPredict(model) +# new implementation with analytical Hessians +gdml = GDMLPredict(model).to(device) +# load geometry +r,_ = io.read_xyz(geometry_file) +coords = torch.from_numpy(r).to(device) + +########################################################################## +# # +# check that energies and gradients agree with reference implementation # +# # +########################################################################## + +# make random numbers reproducible +torch.manual_seed(0) + +natom = coords.size()[1]//3 +# timings for different batch sizes +for batch_size in [1,10,100,1000]: + logger.info(f"batch size {batch_size}") + # batch (B,3*N) + rs = coords.repeat(batch_size, 1) + 0.1 * torch.rand(batch_size,3*natom).to(device) + # (B, N, 3) + rs_3N = rs.reshape(batch_size, -1, 3) + + t_start = time.time() + # compute energy and Hessian with reference implementation + en_ref, force_ref = gdml_ref.forward(rs_3N) + grad_ref = -force_ref.reshape(rs.size()) + + t_end = time.time() + logger.info(f"timing reference implementation, energy+gradient : {t_end-t_start} seconds") + + t_start = time.time() + # and compare with new implementation + en, grad, hessian = gdml.forward(rs) + + t_end = time.time() + logger.info(f"timing new implementation, energy+gradient+hessian : {t_end-t_start} seconds") + + # error per sample + err_en = torch.norm(en_ref - en)/batch_size + err_grad = torch.norm(grad_ref - grad)/batch_size + + logger.info(f" error of energy : {err_en}") + logger.info(f" error of gradient : {err_grad}") + + assert err_en < 1.0e-4 + assert err_grad < 1.0e-4 + +############################################################### +# # +# compare numerical and analytic Hessians of sGDML potential # +# # +############################################################### +from sgdml.intf.ase_calc import SGDMLCalculator +from ase.io.xyz import read_xyz +from ase.optimize import BFGS +from ase.vibrations import Vibrations +from ase.units import kcal, mol + +# compute Hessian numerically using ASE +with open(geometry_file) as f: + molecule = next(read_xyz(f)) +sgdml_calc = SGDMLCalculator(model_file) +molecule.calc = sgdml_calc + +# optimization +opt = BFGS(molecule) +opt.run(fmax=0.001) +# optimized geometry +coords_opt = torch.from_numpy(molecule.get_positions()).reshape(1,-1).to(device) + +# frequencies +vib = Vibrations(molecule, name="/tmp/vib_sgdml") +vib.run() +vib.get_energies() +vib.clean() + +# convert numerical Hessian from eV Ang^{-2} to kcal/mol Ang^{-2} +hessian_numerical = vib.H / (kcal / mol) + + +# compute analytic Hessian directly from sGDML model +hessian_analytical = gdml.forward(coords_opt)[2][0,:,:].cpu().numpy() + +# check that Hessian is symmetric +err_sym = la.norm(hessian_analytical - hessian_analytical.T) +logger.info(f"|Hessian-Hessian^T|= {err_sym}") +assert err_sym < 1.0e-8 + +""" +# compare Hessians visually +import matplotlib.pyplot as plt +ax1 = plt.subplot(1,3,1) +ax1.set_title("numerical Hessian") +ax1.imshow(hessian_numerical) + +ax2 = plt.subplot(1,3,2) +ax2.set_title("analytical Hessian") +ax2.imshow(hessian_analytical) + +ax3 = plt.subplot(1,3,3) +ax3.set_title("difference") +ax3.imshow(hessian_numerical - hessian_analytical) + +plt.show() +""" + +# check that numerical and analytical Hessians agree within numerical errors +err = la.norm(hessian_numerical - hessian_analytical)/la.norm(hessian_numerical) +logger.info(f"|Hessian(num)-Hessian(ana)|/|Hessian(num)|= {err}") +assert err < 1.0e-3 diff --git a/sgdml/torchtools_hessian.py b/sgdml/torchtools_hessian.py new file mode 100644 index 0000000..08dd8ba --- /dev/null +++ b/sgdml/torchtools_hessian.py @@ -0,0 +1,239 @@ +# coding: utf-8 +"""sGDML force field with analytic energies, gradients and Hessians""" + +__all__ = ['GDMLPredict'] + +# MIT License +# +# Copyright (c) 2019-2020 Jan Hermann, Stefan Chmiela +# modified by Alexander Humeniuk +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +import sys +import numpy as np +import torch +import torch.nn as nn + + +class GDMLPredict(nn.Module): + def __init__(self, model): + """ + Predict molecular energies, gradients and Hessians from machine-learned GDML model. + + This is a modified version of Stefan Chmiela's GDML potential (adapted from https://github.com/stefanch/sGDML) + which can also predict second order derivatives (Hessians) of the potential energy. + + Parameters + ---------- + model : Mapping + Obtained from :meth:`~train.GDMLTrain.train`. + It is assumed that the model uses atomic units (bohr for lengths and Hartree for energies). + + Notes + ----- + On a CPU, a Hessian calculation is roughly 5 times more expensive than a gradient calculation. On a GPU, it is 10 times + more expensive than a gradient calculation. Batches of approximately 10000 medium-sized molecules can be computed per second + on a GPU. + """ + super().__init__() + + model = dict(model) + + self._sig = int(model['sig']) + self._c = float(model['c']) + self._std = float(model.get('std', 1)) + + self.n_atoms = model['z'].shape[0] + + desc_siz = model['R_desc'].shape[0] + n_perms, self._n_atoms = model['perms'].shape + perm_idxs = ( + torch.tensor(model['tril_perms_lin']) + .view(-1, n_perms) + .t() + ) + + self._xs_train, self._Jx_alphas = ( + nn.Parameter( + xs.repeat(1, n_perms)[:, perm_idxs].reshape(-1, desc_siz), + requires_grad=False, + ) + for xs in ( + torch.tensor(model['R_desc']).t(), + torch.tensor(np.array(model['R_d_desc_alpha'])), + ) + ) + + self.perm_idxs = perm_idxs + self.n_perms = n_perms + + def forward(self, r, order=2): + """ + Predict + * molecular energy, + * gradient of the energy, + * matrix of second order derivatives (Hessian) of energy + for a batch of geometries. + + Parameters + ---------- + r : Tensor + (dims B x 3N) Cartesian coordinates (in bohr) of M molecules composed of N atoms + + Returns + ------- + energy : Tensor + (dims B) Molecular energies (in Hartree) + grad : Tensor + (dims B x 3N) Gradients of molecular energies (in Hartree * bohr^{-1}) + hess : Tensor + (dims B x 3N x 3N) Hessian matrices (in Hartree * bohr^{-2}) + + What is returned depends on the option `order`. + + + Optional + -------- + order : int + Order of the highest derivative that will be returned (0 - E, 1 - dE/dx, 2 - d^2E/dxdy). + Allows to stop the calculation early, if only the energy and/or gradient + is needed. Depending on the value of `order` 1,2 or 3 Tensors are returned: + order=0 - energy + order=1 - (energy, grad) + order=2 - (energy, grad, hess) + """ + # dimensions + # B: batch size, number of molecular for which predictions should be made + # M: number of training samples + # N: number of atoms + # D: dimension of descriptor, for Coulomb matrix D = N*(N-1)/2 + # X: spatial dimensions, 3 + # + + # dimensions + dimN = self._n_atoms + dimM, dimD = self._Jx_alphas.size() + dimB = r.size()[0] + assert r.size()[1] == 3*dimN + + sig = self._sig + q = np.sqrt(5) / sig + + r = torch.reshape(r, (-1, dimN, 3)) + # + diffs = r[:, :, None, :] - r[:, None, :, :] # (B, N, N, 3) + + dists = diffs.norm(dim=-1) # (B, N, N) + + # indices of lower tiangular matrix, i > j + i,j = np.tril_indices(dimN, k=-1) + xs = 1.0 / dists[:, i, j] # (B, D) + + del dists + + x_diffs = xs[:, None, :] - self._xs_train # (B, M, D) + x_dists = x_diffs.norm(dim=-1) # (B, M) + + A = self._Jx_alphas + # XA = sum_n (x_n - x_n') A_n + XA = torch.einsum('bmd,md->bm', x_diffs, A) # (B, M) + + exp_fac = 1.0/3.0 * q**4 * torch.exp(-q * x_dists) # (B, M) + + energy = torch.einsum('bm,bm->b', exp_fac * (1.0 + q*x_dists)/q**2, XA) + energy = (energy * self._std + self._c) + + if order == 0: + # 0-th order derivative + return energy + + # construct gradient of molecular energy + + xs3 = xs**3 + # chain rule: gradient w/r/t descriptor --> gradient w/r/t cartesian coordinates + + # construct Jacobian of Coulomb matrix D_(ij) = 1/|r(i)-r(j)| + jacobian = torch.zeros(dimB, dimD, dimN, 3, + dtype=diffs.dtype, + device=A.device) # (B, D, N, 3) + k,l = torch.tril_indices(dimN, dimN, offset=-1) + kl = torch.arange(dimD) + jacobian[:,kl,k,:] = -xs3[:,:,None] * diffs[:,k,l,:] + jacobian[:,kl,l,:] -= xs3[:,:,None] * diffs[:,l,k,:] + jacobian = torch.reshape(jacobian, (dimB, dimD, 3*dimN)) + + grad_x = torch.einsum('bm,md->bd', exp_fac * (1.0 + q*x_dists)/q**2, A) # (B, D) + grad_x -= torch.einsum('bm,bmd->bd', exp_fac * XA, x_diffs) + # transform gradient to cartesian coordinates + grad = torch.einsum('bd,bdx->bx', grad_x, jacobian) + grad *= self._std + + if order == 1: + # 0-th and 1st order derivatives + return energy, grad + + # construct Hessian of molecular energy + # XJ = sum_a (x_a - x_a') J_ax + XJ = torch.einsum('bmd,bdx->bmx', x_diffs, jacobian) # (B, M, 3*N) + AJ = torch.einsum('md,bdx->bmx', A, jacobian) # (B, M, 3*N) + JJ = torch.einsum('bdx,bdy->bxy', jacobian, jacobian) # (B, 3*N, 3*N) + + del jacobian + + # sum over training set (M dimension) + hess = torch.einsum('bm,bmx,bmy->bxy', + exp_fac * XA * q/x_dists, XJ, XJ) # (B, 3*N, 3*N) + hess -= torch.einsum('bm,bxy->bxy', exp_fac * XA, JJ) + hess -= torch.einsum('bm,bmx,bmy->bxy', exp_fac, AJ, XJ) + hess -= torch.einsum('bm,bmx,bmy->bxy', exp_fac, XJ, AJ) + + del XA, XJ, AJ, JJ, x_dists, exp_fac + + h1 = ( 3 * grad_x[:,kl,None,None] * xs[:,kl,None,None]**5 + * diffs[:,k,l,:,None] * diffs[:,k,l,None,:] ) + h2 = -grad_x[:,kl] * xs[:,kl]**3 + + idxB = torch.arange(dimB).unsqueeze(1).expand(dimB, dimD) + # loop over cartesian coordinates u,v = x,y,z + for u in [0,1,2]: + for v in [0,1,2]: + h1_uv = h1[:,:,u,v] + hess[:,3*k+u,3*l+v] -= h1_uv + hess[:,3*l+u,3*k+v] -= h1_uv + + # Because the index arrays k and l contain repeated indices + # we cannot simply use the notation + # hess[:,3*k+u,3*k+v] += h1[:,kl,u,v] + # hess[:,3*l+u,3*l+v] += h1[:,kl,u,v] + # since the the contributions from repeated indices are not + # accumulated. Instead we have to use `index_put_(..., accumulate=True)`. + hess.index_put_((idxB,3*k+u,3*k+v), h1[:,kl,u,v], accumulate=True) + hess.index_put_((idxB,3*l+u,3*l+v), h1[:,kl,u,v], accumulate=True) + + hess[:,3*k+u,3*l+u] -= h2 + hess[:,3*l+u,3*k+u] -= h2 + + hess.index_put_((idxB,3*k+u,3*k+u), h2[:,kl], accumulate=True) + hess.index_put_((idxB,3*l+u,3*l+u), h2[:,kl], accumulate=True) + + hess *= self._std + + # 0-th, 1st and 2nd order derivatives + return energy, grad, hess diff --git a/sgdml/train.py b/sgdml/train.py index 3b9dbe8..38f6d6c 100755 --- a/sgdml/train.py +++ b/sgdml/train.py @@ -47,7 +47,7 @@ from . import __version__, DONE, NOT_DONE from .solvers.analytic import Analytic -from .solvers.iterative import Iterative +#from .solvers.iterative import Iterative from .predict import GDMLPredict from .utils.desc import Desc from .utils import io, perm, ui From c9bc981113caeb6ae62044f600943826008fffe4 Mon Sep 17 00:00:00 2001 From: Alexander Humeniuk Date: Sat, 16 Jan 2021 22:07:30 +0100 Subject: [PATCH 2/9] When the comment line in an xyz file contains additional data fields, the energy label is read as the first float in the line --- scripts/sgdml_dataset_from_extxyz.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/scripts/sgdml_dataset_from_extxyz.py b/scripts/sgdml_dataset_from_extxyz.py index b599377..2afe4d9 100755 --- a/scripts/sgdml_dataset_from_extxyz.py +++ b/scripts/sgdml_dataset_from_extxyz.py @@ -57,7 +57,8 @@ def read_nonstd_ext_xyz(f): if line_i == 1: try: - e = float(line) + parts = line.split() + e = float(parts[0]) except ValueError: pass else: From 57cc76a16be0dcea13e2f6914bf6b66de049e5bd Mon Sep 17 00:00:00 2001 From: Alexander Humeniuk Date: Sun, 17 Jan 2021 01:57:07 +0100 Subject: [PATCH 3/9] fix --- scripts/sgdml_dataset_from_extxyz.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/scripts/sgdml_dataset_from_extxyz.py b/scripts/sgdml_dataset_from_extxyz.py index 2afe4d9..b599377 100755 --- a/scripts/sgdml_dataset_from_extxyz.py +++ b/scripts/sgdml_dataset_from_extxyz.py @@ -57,8 +57,7 @@ def read_nonstd_ext_xyz(f): if line_i == 1: try: - parts = line.split() - e = float(parts[0]) + e = float(line) except ValueError: pass else: From 9e0fa025d2cc5f0dbb6f3ad12e5b05a8ef7dc94c Mon Sep 17 00:00:00 2001 From: Alexander Humeniuk Date: Thu, 21 Jan 2021 16:33:37 +0100 Subject: [PATCH 4/9] Energy labels can be read via Energy=... from the comment line in an ordinary xyz file --- scripts/sgdml_dataset_from_extxyz.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/scripts/sgdml_dataset_from_extxyz.py b/scripts/sgdml_dataset_from_extxyz.py index b599377..6d54257 100755 --- a/scripts/sgdml_dataset_from_extxyz.py +++ b/scripts/sgdml_dataset_from_extxyz.py @@ -27,6 +27,7 @@ import argparse import os import sys +import re try: from ase.io import read @@ -45,7 +46,8 @@ # Assumes that the atoms in each molecule are in the same order. def read_nonstd_ext_xyz(f): n_atoms = None - + pattern = re.compile('[eE]nergy=([\+\-0-9\.]+) ') + R, z, E, F = [], [], [], [] for i, line in enumerate(f): line = line.strip() @@ -59,7 +61,12 @@ def read_nonstd_ext_xyz(f): try: e = float(line) except ValueError: - pass + # Try to read energy from comment line as + # Energy=(.*) ... + match = pattern.findall(line) + if len(match) > 0: + e = float(match[0]) + E.append(e) else: E.append(e) From 8ec826e96cfb4454878e06fd3bc063873d8c188a Mon Sep 17 00:00:00 2001 From: Alexander Humeniuk Date: Tue, 26 Jan 2021 10:43:40 +0100 Subject: [PATCH 5/9] store level of theory in .npz file --- scripts/sgdml_dataset_from_extxyz.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/scripts/sgdml_dataset_from_extxyz.py b/scripts/sgdml_dataset_from_extxyz.py index 6d54257..a88a386 100755 --- a/scripts/sgdml_dataset_from_extxyz.py +++ b/scripts/sgdml_dataset_from_extxyz.py @@ -196,6 +196,16 @@ def read_nonstd_ext_xyz(f): base_vars['F_min'], base_vars['F_max'] = np.min(F.ravel()), np.max(F.ravel()) base_vars['F_mean'], base_vars['F_var'] = np.mean(F.ravel()), np.var(F.ravel()) +print('Please provide a name for this dataset. Otherwise the original filename will be reused.') +custom_name = raw_input('> ').strip() +if custom_name != '': + name = custom_name + +print('Please provide a descriptor for the level of theory used to create this dataset.') +theory = raw_input('> ').strip() +if theory == '': + theory = 'unknown' + print('Please provide a description of the length unit used in your input file, e.g. \'Ang\' or \'au\': ') print('Note: This string will be stored in the dataset file and passed on to models files for later reference.') r_unit = raw_input('> ').strip() From e7fa121371351c02c68c877ac1fb41371a8b1b8d Mon Sep 17 00:00:00 2001 From: Alexander Humeniuk Date: Tue, 26 Jan 2021 10:54:11 +0100 Subject: [PATCH 6/9] level of theory --- scripts/sgdml_dataset_from_extxyz.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/scripts/sgdml_dataset_from_extxyz.py b/scripts/sgdml_dataset_from_extxyz.py index a88a386..670d43a 100755 --- a/scripts/sgdml_dataset_from_extxyz.py +++ b/scripts/sgdml_dataset_from_extxyz.py @@ -196,15 +196,11 @@ def read_nonstd_ext_xyz(f): base_vars['F_min'], base_vars['F_max'] = np.min(F.ravel()), np.max(F.ravel()) base_vars['F_mean'], base_vars['F_var'] = np.mean(F.ravel()), np.var(F.ravel()) -print('Please provide a name for this dataset. Otherwise the original filename will be reused.') -custom_name = raw_input('> ').strip() -if custom_name != '': - name = custom_name - print('Please provide a descriptor for the level of theory used to create this dataset.') theory = raw_input('> ').strip() if theory == '': theory = 'unknown' +base_vars['theory'] = theory print('Please provide a description of the length unit used in your input file, e.g. \'Ang\' or \'au\': ') print('Note: This string will be stored in the dataset file and passed on to models files for later reference.') From 8e997fbb286311ad37fa5204479efa38a681f8ee Mon Sep 17 00:00:00 2001 From: Alexander Humeniuk Date: Sun, 7 Feb 2021 14:47:51 +0100 Subject: [PATCH 7/9] index arrays were inadvertently forgotten on the CPU, after moving all index arrays to the GPU a 10x speed-up is achieved --- sgdml/torchtools_hessian.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sgdml/torchtools_hessian.py b/sgdml/torchtools_hessian.py index 08dd8ba..5167083 100644 --- a/sgdml/torchtools_hessian.py +++ b/sgdml/torchtools_hessian.py @@ -173,8 +173,8 @@ def forward(self, r, order=2): jacobian = torch.zeros(dimB, dimD, dimN, 3, dtype=diffs.dtype, device=A.device) # (B, D, N, 3) - k,l = torch.tril_indices(dimN, dimN, offset=-1) - kl = torch.arange(dimD) + k,l = torch.tril_indices(dimN, dimN, offset=-1, device=A.device) + kl = torch.arange(dimD, device=A.device) jacobian[:,kl,k,:] = -xs3[:,:,None] * diffs[:,k,l,:] jacobian[:,kl,l,:] -= xs3[:,:,None] * diffs[:,l,k,:] jacobian = torch.reshape(jacobian, (dimB, dimD, 3*dimN)) From ca084b69dd74ede6919d3bfd3037f82befb351ac Mon Sep 17 00:00:00 2001 From: Alexander Humeniuk Date: Sun, 7 Feb 2021 14:52:35 +0100 Subject: [PATCH 8/9] index arrays were inadvertently forgotten on the CPU, after moving all index arrays to the GPU a 10x speed-up is achieved --- sgdml/torchtools_hessian.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/sgdml/torchtools_hessian.py b/sgdml/torchtools_hessian.py index 5167083..5833766 100644 --- a/sgdml/torchtools_hessian.py +++ b/sgdml/torchtools_hessian.py @@ -84,6 +84,15 @@ def __init__(self, model): self.perm_idxs = perm_idxs self.n_perms = n_perms + @property + def device(self): + """ + device (gpu or cuda) where the GDML model lives on + """ + # nn.Module does not have a .device attribute, so we take the value + # from one of the tensors. + return self._xs_train.device + def forward(self, r, order=2): """ Predict @@ -132,7 +141,8 @@ def forward(self, r, order=2): dimM, dimD = self._Jx_alphas.size() dimB = r.size()[0] assert r.size()[1] == 3*dimN - + device=self.device + sig = self._sig q = np.sqrt(5) / sig @@ -143,7 +153,7 @@ def forward(self, r, order=2): dists = diffs.norm(dim=-1) # (B, N, N) # indices of lower tiangular matrix, i > j - i,j = np.tril_indices(dimN, k=-1) + i,j = np.tril_indices(dimN, k=-1, device=device) xs = 1.0 / dists[:, i, j] # (B, D) del dists @@ -172,9 +182,9 @@ def forward(self, r, order=2): # construct Jacobian of Coulomb matrix D_(ij) = 1/|r(i)-r(j)| jacobian = torch.zeros(dimB, dimD, dimN, 3, dtype=diffs.dtype, - device=A.device) # (B, D, N, 3) - k,l = torch.tril_indices(dimN, dimN, offset=-1, device=A.device) - kl = torch.arange(dimD, device=A.device) + device=device) # (B, D, N, 3) + k,l = torch.tril_indices(dimN, dimN, offset=-1, device=device) + kl = torch.arange(dimD, device=device) jacobian[:,kl,k,:] = -xs3[:,:,None] * diffs[:,k,l,:] jacobian[:,kl,l,:] -= xs3[:,:,None] * diffs[:,l,k,:] jacobian = torch.reshape(jacobian, (dimB, dimD, 3*dimN)) @@ -210,7 +220,7 @@ def forward(self, r, order=2): * diffs[:,k,l,:,None] * diffs[:,k,l,None,:] ) h2 = -grad_x[:,kl] * xs[:,kl]**3 - idxB = torch.arange(dimB).unsqueeze(1).expand(dimB, dimD) + idxB = torch.arange(dimB, device=device).unsqueeze(1).expand(dimB, dimD) # loop over cartesian coordinates u,v = x,y,z for u in [0,1,2]: for v in [0,1,2]: From b794c394e1dff50c783b2980aa22dde8cc3e035b Mon Sep 17 00:00:00 2001 From: Alexander Humeniuk Date: Sun, 7 Feb 2021 14:58:05 +0100 Subject: [PATCH 9/9] after moving all indices to GPU, Hessian calculation is only 2x to 4x more expensive than a gradient calculation --- sgdml/_bmark_cache.npz | Bin 0 -> 1207 bytes sgdml/torchtools_hessian.py | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) create mode 100644 sgdml/_bmark_cache.npz diff --git a/sgdml/_bmark_cache.npz b/sgdml/_bmark_cache.npz new file mode 100644 index 0000000000000000000000000000000000000000..fc30fe4970b66a23b5dd768c90f10b7d3ebfb103 GIT binary patch literal 1207 zcmWIWW@Zs#U|`??VnqhMO+nitfvgokEWjYbker{A8ef)LRGgWgrRBmPQS;N*#bV*zXDO^TWmvm|fnf;{1@j0Ni!FPJwiJ|t+$4DtE#QzmguD)2s| z>Brh7o~|Ocgz2Z$EGw?9TvHhInAzCaY!%EJFNpuY#^a8khKtdn@k_G|n?wOpn5Iuy1%27iI>KqY j - i,j = np.tril_indices(dimN, k=-1, device=device) + i,j = torch.tril_indices(dimN, dimN, offset=-1, device=device) xs = 1.0 / dists[:, i, j] # (B, D) del dists