Skip to content

Commit

Permalink
Initial support for using msgpack coordinates / energies
Browse files Browse the repository at this point in the history 
Signed-off-by: Geoff Hutchison <[email protected]>
  • Loading branch information
@ghutchis
ghutchis committed Nov 20, 2024
1 parent 86eeff4 commit 23447e4
Show file tree
Hide file tree
Showing 9 changed files with 335 additions and 103 deletions.
130 changes: 93 additions & 37 deletions avogadro/qtplugins/forcefield/scriptenergy.cpp
View file
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()

Check warning on line 64 in avogadro/qtplugins/forcefield/scriptenergy.cpp

View check run for this annotation

Codacy Production / Codacy Static Code Analysis

avogadro/qtplugins/forcefield/scriptenergy.cpp#L64 

Either the condition 'mol==nullptr' is redundant or there is possible null pointer dereference: mol.
<< " method: " << m_identifier.c_str()
<< " script: " << m_interpreter->scriptFilePath() << " msgPack "

Check warning on line 66 in avogadro/qtplugins/forcefield/scriptenergy.cpp

View check run for this annotation

Codacy Production / Codacy Static Code Analysis

avogadro/qtplugins/forcefield/scriptenergy.cpp#L66 

Either the condition 'm_interpreter==nullptr' is redundant or there is possible null pointer dereference: m_interpreter.
<< 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
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
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
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

0 comments on commit 23447e4

Please sign in to comment.