Node featurization does not properly assign hybridization using OpenEye toolkit

espaloma: Extensible Surrogate Potential Optimized by Message-passing Algorithms 🍹

Source code for Wang Y, Fass J, and Chodera JD "End-to-End Differentiable Construction of Molecular Mechanics Force Fields."

Documentation: https://docs.espaloma.org

Paper Abstract

Molecular mechanics (MM) potentials have long been a workhorse of computational chemistry. Leveraging accuracy and speed, these functional forms find use in a wide variety of applications in biomolecular modeling and drug discovery, from rapid virtual screening to detailed free energy calculations. Traditionally, MM potentials have relied on human-curated, inflexible, and poorly extensible discrete chemical perception rules atom types for applying parameters to small molecules or biopolymers, making it difficult to optimize both types and parameters to fit quantum chemical or physical property data. Here, we propose an alternative approach that uses graph neural networks to perceive chemical environments, producing continuous atom embeddings from which valence and nonbonded parameters can be predicted using invariance-preserving layers. Since all stages are built from smooth neural functions, the entire process---spanning chemical perception to parameter assignment---is modular and end-to-end differentiable with respect to model parameters, allowing new force fields to be easily constructed, extended, and applied to arbitrary molecules. We show that this approach is not only sufficiently expressive to reproduce legacy atom types, but that it can learn and extend existing molecular mechanics force fields, construct entirely new force fields applicable to both biopolymers and small molecules from quantum chemical calculations, and even learn to accurately predict free energies from experimental observables.

Installation

$ conda install -c conda-forge "espaloma=0.3.2"

Example: Deploy espaloma 0.3.2 pretrained force field to arbitrary MM system

# imports
import os
import torch
import espaloma as esp

# define or load a molecule of interest via the Open Force Field toolkit
from openff.toolkit.topology import Molecule
molecule = Molecule.from_smiles("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")

# create an Espaloma Graph object to represent the molecule of interest
molecule_graph = esp.Graph(molecule)

# load pretrained model
espaloma_model = esp.get_model("latest")

# apply a trained espaloma model to assign parameters
espaloma_model(molecule_graph.heterograph)

# create an OpenMM System for the specified molecule
openmm_system = esp.graphs.deploy.openmm_system_from_graph(molecule_graph)

If using espaloma from a local .pt file, say for example espaloma-0.3.2.pt, then you would need to run the eval method of the model to get the correct inference/predictions, as follows:

import torch
...
# load local pretrained model
espaloma_model = torch.load("espaloma-0.3.2.pt")
espaloma_model.eval()
...

The rest of the code should be the same as in the previous code block example.

Compatible models

Below is a compatibility matrix for different versions of espaloma code and espaloma models (the .pt file).

Model 🧪	Supported Espaloma version 💻	Release Date 🗓️	Espaloma architecture change 📐?
`espaloma-0.3.2.pt`	0.3.1, 0.3.2	Sep 22, 2023	✅ No
`espaloma-0.3.1.pt`	0.3.1, 0.3.2	Jul 17, 2023	⚠️ Yes
`espaloma-0.3.0.pt`	0.3.0	Apr 26, 2023	⚠️Yes

Note

espaloma-0.3.1.pt and espaloma-0.3.2.pt are the same model.

Using espaloma to parameterize small molecules in relative free energy calculations

An example of using espaloma to parameterize small molecules in relative alchemical free energy calculations is provided in the scripts/perses-benchmark/ directory.

Manifest

espaloma/ core code for graph-parametrized potential energy functions.
- graphs/ data objects that contain various level of information we need.
  - graph.py base modules for graphs.
  - molecule_graph.py provide APIs to various molecular modelling toolkits.
  - homogeneous_graph.py simplest graph representation of a molecule.
  - heterogeneous_graph.py graph representation of a molecule that contains information regarding membership of lower-level nodes to higher-level nodes.
  - parametrized_graph.py graph representation of a molecule with all parameters needed for energy evaluation.
- nn/ neural network models that facilitates translation between graphs.
  - dgl_legacy.py API to dgl models for atom-level message passing.
- mm/ molecular mechanics functionalities for energy evaluation.
  - i/ energy terms used in Class-I force field.
    - bond.py bond energy
    - angle.py angle energy
    - torsion.py torsion energy
    - nonbonded.py nonbonded energy
  - ii/ energy terms used in Class-II force field.
    - coupling.py coupling terms
    - polynomial.py higher order polynomials.

License

This software is licensed under MIT license.

Copyright

Yuanqing Wang
Josh Fass
John D. Chodera

	def dihedral(
	x0: torch.Tensor, x1: torch.Tensor, x2: torch.Tensor, x3: torch.Tensor
	) -> torch.Tensor:
	""" Dihedral between four points.

	Reference
	---------
	Closely follows implementation in Yutong Zhao's timemachine:
	https://github.com/proteneer/timemachine/blob/1a0ab45e605dc1e28c44ea90f38cb0dedce5c4db/timemachine/potentials/bonded.py#L152-L199
	"""
	# check input shapes

	assert x0.shape == x1.shape == x2.shape == x3.shape

	# compute displacements 0->1, 2->1, 2->3
	r01 = x1 - x0
	r21 = x1 - x2
	r23 = x3 - x2

	# compute normal planes
	n1 = torch.cross(r01, r21)
	n2 = torch.cross(r21, r23)

	rkj_normed = r21 / torch.norm(r21, dim=-1, keepdim=True)

	y = torch.sum(torch.mul(torch.cross(n1, n2), rkj_normed), dim=-1)
	x = torch.sum(torch.mul(n1, n2), dim=-1)

	# choose quadrant correctly
	theta = torch.atan2(y, x)

	return theta

	def periodic(
	x, k, periodicity=list(range(1, 7)), phases=[0.0 for _ in range(6)]
	):
	""" Periodic term.

	Parameters
	----------
	x : `torch.Tensor`, `shape=(batch_size, 1)`
	k : `torch.Tensor`, `shape=(batch_size, number_of_phases)`
	periodicity: either list of length number_of_phases, or
	`torch.Tensor`, `shape=(batch_size, number_of_phases)`
	phases : either list of length number_of_phases, or
	`torch.Tensor`, `shape=(batch_size, number_of_phases)`
	"""

	if isinstance(phases, list):
	phases = torch.tensor(phases, device=x.device)

	if isinstance(periodicity, list):
	periodicity = torch.tensor(
	periodicity, device=x.device, dtype=torch.get_default_dtype(),
	)

	if periodicity.ndim == 1:
	periodicity = periodicity[None, None, :].repeat(
	x.shape[0], x.shape[1], 1
	)

	elif periodicity.ndim == 2:
	periodicity = periodicity[:, None, :].repeat(1, x.shape[1], 1)

	if phases.ndim == 1:
	phases = phases[None, None, :].repeat(x.shape[0], x.shape[1], 1,)

	elif phases.ndim == 2:
	phases = phases[:, None, :].repeat(1, x.shape[1], 1,)

	n_theta = periodicity * x[:, :, None]

	n_theta_minus_phases = n_theta - phases

	cos_n_theta_minus_phases = n_theta_minus_phases.cos()

	k = k[:, None, :].repeat(1, x.shape[1], 1)

	energy = (k * (1.0 + cos_n_theta_minus_phases)).sum(dim=-1)

	return energy

	def apply_torsion(node, n_max_phases=6):
	phases = torch.zeros(
	g.heterograph.number_of_nodes("n4"), n_max_phases,
	)

	periodicity = torch.zeros(
	g.heterograph.number_of_nodes("n4"), n_max_phases,
	)

	k = torch.zeros(g.heterograph.number_of_nodes("n4"), n_max_phases,)

	force = forces["ProperTorsions"]

	for idx in range(g.heterograph.number_of_nodes("n4")):
	idxs = tuple(node.data["idxs"][idx].numpy())
	if idxs in force:
	_force = force[idxs]
	for sub_idx in range(len(_force.periodicity)):
	if hasattr(_force, "k%s" % sub_idx):
	k[idx, sub_idx] = getattr(
	_force, "k%s" % sub_idx
	).value_in_unit(esp.units.ENERGY_UNIT)

	phases[idx, sub_idx] = getattr(
	_force, "phase%s" % sub_idx
	).value_in_unit(esp.units.ANGLE_UNIT)

	periodicity[idx, sub_idx] = getattr(
	_force, "periodicity%s" % sub_idx
	)

	return {
	"k_ref": k,
	"periodicity_ref": periodicity,
	"phases_ref": phases,
	}

	g.heterograph.apply_nodes(apply_torsion, ntype="n4")

choderalab / espaloma Goto Github PK

espaloma's Introduction

espaloma: Extensible Surrogate Potential Optimized by Message-passing Algorithms 🍹

Paper Abstract

Installation

Example: Deploy espaloma 0.3.2 pretrained force field to arbitrary MM system

Compatible models

Using espaloma to parameterize small molecules in relative free energy calculations

Manifest

License

Copyright

espaloma's People

Contributors

Stargazers

Watchers

Forkers

espaloma's Issues

Problem

Reproduce (atom hybridization):

Problem

Recommend Projects

Recommend Topics

Recommend Org