Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optionally use msgpack for energy / gradient calculations #1804

Closed
wants to merge 3 commits into from
Closed

Optionally use msgpack for energy / gradient calculations #1804

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
86eeff4
Fix some minor bugs with vector support in the Variant class
ghutchis Nov 20, 2024
23447e4
Initial support for using msgpack coordinates / energies
ghutchis Nov 20, 2024
3f72754
Merge branch 'master' into msgpack-calcs
ghutchis Nov 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 93 additions & 37 deletions avogadro/qtplugins/forcefield/scriptenergy.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,10 @@
#include <qjsonobject.h>
#include <qjsonvalue.h>

// nlohmann json for msgpack
#include <nlohmann/json.hpp>
using json = nlohmann::json;

namespace Avogadro::QtPlugins {

ScriptEnergy::ScriptEnergy(const QString& scriptFileName_)
Expand Down Expand Up @@ -57,6 +61,11 @@ void ScriptEnergy::setMolecule(Core::Molecule* mol)
{
m_molecule = mol;

qDebug() << " ScriptEnergy::setMolecule() " << mol->formula().c_str()
<< " method: " << m_identifier.c_str()
<< " script: " << m_interpreter->scriptFilePath() << " msgPack "
<< m_msgPack;

// should check if the molecule is valid for this script
// .. this should never happen, but let's be defensive
if (mol == nullptr || m_interpreter == nullptr) {
Expand Down Expand Up @@ -116,23 +125,41 @@ Real ScriptEnergy::value(const Eigen::VectorXd& x)

// write the new coordinates and read the energy
QByteArray input;
for (Index i = 0; i < x.size(); i += 3) {
// write as x y z (space separated)
input += QString::number(x[i]).toUtf8() + " " +
QString::number(x[i + 1]).toUtf8() + " " +
QString::number(x[i + 2]).toUtf8() + "\n";
if (m_msgPack) {
json j;
j["coordinates"] = json::array();
for (Index i = 0; i < x.size(); ++i) {
j["coordinates"].push_back(x[i]);
}
auto msgPack = json::to_msgpack(j);
input = QByteArray::fromRawData(
reinterpret_cast<const char*>(msgPack.data()), msgPack.size());
} else {
for (Index i = 0; i < x.size(); i += 3) {
// write as x y z (space separated)
input += QString::number(x[i]).toUtf8() + " " +
QString::number(x[i + 1]).toUtf8() + " " +
QString::number(x[i + 2]).toUtf8() + "\n";
}
}
QByteArray result = m_interpreter->asyncWriteAndResponse(input);

// go through lines in result until we see "AvogadroEnergy: "
QStringList lines = QString(result).remove('\r').split('\n');
double energy = 0.0;
for (auto line : lines) {
if (line.startsWith("AvogadroEnergy:")) {
QStringList items = line.split(" ", Qt::SkipEmptyParts);
if (items.size() > 1) {
energy = items[1].toDouble();
break;
if (m_msgPack) {
json j = json::from_msgpack(result.toStdString());
if (j.find("energy") != j.end()) {
return j["energy"];
}
} else {
// go through lines in result until we see "AvogadroEnergy: "
QStringList lines = QString(result).remove('\r').split('\n');
for (auto line : lines) {
if (line.startsWith("AvogadroEnergy:")) {
QStringList items = line.split(" ", Qt::SkipEmptyParts);
if (items.size() > 1) {
energy = items[1].toDouble();
break;
}
}
}
}
Expand All @@ -150,36 +177,56 @@ void ScriptEnergy::gradient(const Eigen::VectorXd& x, Eigen::VectorXd& grad)
// Get the gradient from the script
// write the new coordinates and read the energy
QByteArray input;
for (Index i = 0; i < x.size(); i += 3) {
// write as x y z (space separated)
input += QString::number(x[i]).toUtf8() + " " +
QString::number(x[i + 1]).toUtf8() + " " +
QString::number(x[i + 2]).toUtf8() + "\n";
if (m_msgPack) {
json j;
j["coordinates"] = json::array();
for (Index i = 0; i < x.size(); ++i) {
j["coordinates"].push_back(x[i]);
}
auto msgPack = json::to_msgpack(j);
input = QByteArray::fromRawData(
reinterpret_cast<const char*>(msgPack.data()), msgPack.size());
} else {
for (Index i = 0; i < x.size(); i += 3) {
// write as x y z (space separated)
input += QString::number(x[i]).toUtf8() + " " +
QString::number(x[i + 1]).toUtf8() + " " +
QString::number(x[i + 2]).toUtf8() + "\n";
}
}
QByteArray result = m_interpreter->asyncWriteAndResponse(input);

// parse the result
// first split on newlines
QStringList lines = QString(result).remove('\r').split('\n');
unsigned int i = 0;
bool readingGrad = false;
for (auto line : lines) {
if (line.startsWith("AvogadroGradient:")) {
readingGrad = true;
continue; // next line
if (m_msgPack) {
json j = json::from_msgpack(result.toStdString());
if (j.find("gradient") != j.end()) {
for (Index i = 0; i < x.size(); ++i) {
grad[i] = j["gradient"][i];
}
}

if (readingGrad) {
QStringList items = line.split(" ", Qt::SkipEmptyParts);
if (items.size() == 3) {
grad[i] = items[0].toDouble();
grad[i + 1] = items[1].toDouble();
grad[i + 2] = items[2].toDouble();
i += 3;
} else {
// parse the result
// first split on newlines
QStringList lines = QString(result).remove('\r').split('\n');
unsigned int i = 0;
bool readingGrad = false;
for (auto line : lines) {
if (line.startsWith("AvogadroGradient:")) {
readingGrad = true;
continue; // next line
}

if (i > x.size())
break;
if (readingGrad) {
QStringList items = line.split(" ", Qt::SkipEmptyParts);
if (items.size() == 3) {
grad[i] = items[0].toDouble();
grad[i + 1] = items[1].toDouble();
grad[i + 2] = items[2].toDouble();
i += 3;
}

if (i > x.size())
break;
}
}
}

Expand Down Expand Up @@ -229,6 +276,7 @@ void ScriptEnergy::resetMetaData()
m_valid = false;
m_gradients = false;
m_ions = false;
m_msgPack = false;
m_radicals = false;
m_unitCells = false;
m_inputFormat = NotUsed;
Expand Down Expand Up @@ -351,6 +399,14 @@ void ScriptEnergy::readMetaData()
}
m_radicals = metaData["radical"].toBool();

// msgpack is optional
// defaults to not used
if (metaData["msgpack"].isBool()) {
m_msgPack = metaData["msgpack"].toBool();
} else if (metaData["msgPack"].isBool()) {
m_msgPack = metaData["msgPack"].toBool();
}

// get the element mask
// (if it doesn't exist, the default is no elements anyway)
m_valid = parseElements(metaData);
Expand Down
1 change: 1 addition & 0 deletions avogadro/qtplugins/forcefield/scriptenergy.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ class ScriptEnergy : public Avogadro::Calc::EnergyCalculator
bool m_ions;
bool m_radicals;
bool m_unitCells;
bool m_msgPack = false;

std::string m_identifier;
std::string m_name;
Expand Down
36 changes: 30 additions & 6 deletions avogadro/qtplugins/forcefield/scripts/ani2x.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,12 @@
import json
import sys

try:
import msgpack
msgpack_available = True
except ImportError:
msgpack_available = False

try:
import torch
import torchani
Expand Down Expand Up @@ -34,6 +40,8 @@ def getMetaData():
"ion": False,
"radical": False,
}
if (msgpack_available):
metaData["msgpack"] = True
return metaData


Expand All @@ -54,19 +62,35 @@ def run(filename):
num_atoms = len(atoms)
while True:
# read new coordinates from stdin
for i in range(num_atoms):
np_coords[i] = np.fromstring(input(), sep=" ")
if (msgpack_available):
# unpack the coordinates
data = msgpack.unpackb(sys.stdin.buffer.read())
np_coords = np.array(data["coordinates"], dtype=float).reshape(-1, 3)
else:
for i in range(num_atoms):
np_coords[i] = np.fromstring(input(), sep=" ")
coordinates = torch.tensor([np_coords], requires_grad=True, device=device)

# first print the energy of these coordinates
energy = model((species, coordinates)).energies
print("AvogadroEnergy:", energy) # in Hartree
if (msgpack_available):
response = { "energy": energy }
else:
print("AvogadroEnergy:", energy) # in Hartree

# now print the gradient on each atom
print("AvogadroGradient:")
derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
for i in range(num_atoms):
print(derivative[0][i][0].item(), derivative[0][i][1].item(), derivative[0][i][2].item())

if (msgpack_available):
gradient = []
for i in range(num_atoms):
gradient.append([derivative[0][i][0].item(), derivative[0][i][1].item(), derivative[0][i][2].item()])
response["gradient"] = gradient
print(msgpack.packb(response))
else:
print("AvogadroGradient:")
for i in range(num_atoms):
print(derivative[0][i][0].item(), derivative[0][i][1].item(), derivative[0][i][2].item())


if __name__ == "__main__":
Expand Down
46 changes: 36 additions & 10 deletions avogadro/qtplugins/forcefield/scripts/gaff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,12 @@
import json
import sys

try:
import msgpack
msgpack_available = True
except ImportError:
msgpack_available = False

try:
from openbabel import pybel
import numpy as np
Expand All @@ -30,6 +36,8 @@ def getMetaData():
"ion": False,
"radical": False,
}
if (msgpack_available):
metaData["msgpack"] = True
return metaData


Expand All @@ -48,24 +56,42 @@ def run(filename):
num_atoms = len(mol.atoms)
while True:
# read new coordinates from stdin
for i in range(num_atoms):
coordinates = np.fromstring(input(), sep=" ")
atom = mol.atoms[i]
atom.OBAtom.SetVector(coordinates[0], coordinates[1], coordinates[2])
if (msgpack_available):
# unpack the coordinates
data = msgpack.unpackb(sys.stdin.buffer.read())
np_coords = np.array(data["coordinates"], dtype=float).reshape(-1, 3)
for i in range(num_atoms):
atom = mol.atoms[i]
atom.OBAtom.SetVector(np_coords[i][0], np_coords[i][1], np_coords[i][2])
else:
for i in range(num_atoms):
coordinates = np.fromstring(input(), sep=" ")
atom = mol.atoms[i]
atom.OBAtom.SetVector(coordinates[0], coordinates[1], coordinates[2])

# update the molecule geometry for the next energy
ff.SetCoordinates(mol.OBMol)

# first print the energy of these coordinates
energy = ff.Energy(True) # in kJ/mol
print("AvogadroEnergy:", energy) # in kJ/mol
if (msgpack_available):
response = {"energy": energy}
else:
print("AvogadroEnergy:", energy) # in kJ/mol

# now print the gradient on each atom
print("AvogadroGradient:")
for atom in mol.atoms:
grad = ff.GetGradient(atom.OBAtom)
print(-1.0*grad.GetX(), -1.0*grad.GetY(), -1.0*grad.GetZ())

if (msgpack_available):
gradient = []
for atom in mol.atoms:
grad = ff.GetGradient(atom.OBAtom)
gradient.append([-1.0*grad.GetX(), -1.0*grad.GetY(), -1.0*grad.GetZ()])
response["gradient"] = gradient
print(msgpack.packb(response))
else:
print("AvogadroGradient:")
for atom in mol.atoms:
grad = ff.GetGradient(atom.OBAtom)
print(-1.0*grad.GetX(), -1.0*grad.GetY(), -1.0*grad.GetZ())


if __name__ == "__main__":
Expand Down
Loading
Loading