DOCSEARCHbyAlgolia
Submit Docs

Espaloma

repository·main·Indexed 42 days ago

https://github.com/choderalab/espaloma

An extensible surrogate potential optimized by message-passing algorithms. Espaloma uses graph neural networks to perceive chemical environments and predict continuous atom embeddings, enabling the end-to-end differentiable construction of molecular mechanics force fields.

Tokens
29.9K
Snippets
161
Records
200
Agent score
53%

What's inside Espaloma

  4. Use espaloma.mm.geometry for molecular geometry calculations

    main

    The espaloma.mm.geometry module provides functions and classes for calculating and managing molecular geometry features, such as bond angles, dihedrals, and distances, typically used in the context of preparing molecular systems for OpenMM or machine learning models.

    Key functional capabilities include:

    • Bond and Angle calculations: apply_bond, angle, and apply_angle.
    • Torsional/Dihedral calculations: apply_torsion and dihedral.
    • Distance measurements: distance.
    • Graph-based geometry: geometry_in_graph and the GeometryInGraph class, which likely manages geometric features within a graph representation of a molecule.
    • Utility functions: copy_src and reduce_stack.

  5. Explore the espaloma.mm module hierarchy

    main

    The espaloma.mm module provides the core functionality for integrating Espaloma-derived force field parameters into OpenMM systems. It is organized into several submodules that handle specific components of the molecular mechanics potential:

    • espaloma.mm.angle: Handles angle-related parameters.
    • espaloma.mm.bond: Handles bond-related parameters.
    • espaloma.mm.energy: Manages energy calculations and terms.
    • espaloma.mm.functional: Contains functional components for the force field.
    • espaloma.mm.geometry: Manages geometric representations and constraints.
    • espaloma.mm.nonbonded: Handles non-bonded interactions (e.g., van der Waals, electrostatic).
    • espaloma.mm.torsion: Handles torsion/dihedral parameters.

  6. Define Espaloma model stages

    main

    An Espaloma model is typically composed of several stages that transform graph data into physical parameters and energies:

    1. Representation Stage: Uses graph neural networks (e.g., esp.nn.layers.dgl_legacy.gn) to create atom latent representations.
    2. Readout Stage: Uses pooling (e.g., esp.nn.readout.janossy.JanossyPooling) to map latent representations to modular MM parameters (atoms, bonds, angles, torsions).
    3. Parameter Transformation: Modules like esp.nn.readout.janossy.ExpCoefficients() transform raw outputs into physical coefficients.
    4. Geometry & Energy: Modules like esp.mm.geometry.GeometryInGraph() and esp.mm.energy.EnergyInGraph() compute the final physical properties.
    5. Charge Equilibrium: esp.nn.readout.charge_equilibrium.ChargeEquilibrium() handles charge assignment.

  7. Define an Espaloma Model for QM Fitting

    main

    An Espaloma model is composed of three main stages:

    1. Representation Stage: Maps the molecular graph to atom latent representations using a Graph Neural Network (e.g., SAGEConv).
    2. Readout Stage: Maps atom latent representations to specific molecular mechanics (MM) parameters (bonds, angles, torsions) using a pooling mechanism like JanossyPooling.
    3. Energy/Geometry Stage: Computes the energy and geometry from the parameters.

    In the readout stage, you define the modular MM parameters the model will assign, such as atom hardness, electronegativity, and coefficients for bonds, angles, and torsions.

    # Stage I: Graph -> Atom Latent Representation
representation = esp.nn.Sequential(
    layer=esp.nn.layers.dgl_legacy.gn("SAGEConv"),
    config=[128, "relu", 128, "relu", 128, "relu"],
)

# Stage II & III: Latent -> MM Parameters -> Energy
readout = esp.nn.readout.janossy.JanossyPooling(
    in_features=128, 
    config=[128, "relu", 128, "relu", 128, "relu"],
    out_features={
        1: {"e": 1, "s": 1},           # atom hardness and electronegativity
        2: {"log_coefficients": 2},    # bond linear combination
        3: {"log_coefficients": 2},    # angle linear combination
        4: {"k": 6},                   # torsion barrier heights
    },
)

espaloma_model = torch.nn.Sequential(
                 representation, 
                 readout, 
                 esp.nn.readout.janossy.ExpCoefficients(),
                 esp.mm.geometry.GeometryInGraph(), 
                 esp.mm.energy.EnergyInGraph(),
)

  8. Understand Espaloma versioning with Versioneer

    main

    Espaloma uses Versioneer to automatically infer the installed version based on git tags and commit distance. The versioning follows PEP 440 and uses the regex \d+.\d+.\d+(?+\d+-[a-z0-9]+).

    • If the current commit matches a git tag, the version is exactly that tag (e.g., espaloma-0.1.2).
    • If the commit is ahead of the last tag, the version is appended with +X (where X is the number of commits ahead) and -YYYYYY (where Y is the git commit hash).

  10. Parametrize an MM system with a pretrained Espaloma force field

    main

    You can deploy a pretrained Espaloma force field to parametrize arbitrary small molecule systems. The workflow involves loading a molecule via the Open Force Field toolkit, converting it into an espaloma.Graph object, loading the model using esp.get_model(), applying the model to the graph's heterograph, and finally generating an OpenMM System using esp.graphs.deploy.openmm_system_from_graph().

    import os
import torch
import espaloma as esp
from openff.toolkit.topology import Molecule

# define or load a molecule of interest via the Open Force Field toolkit
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)

  11. Run Espaloma benchmarks via CLI

    main

    The run_benchmarks.py script allows you to run automated benchmarks from the OpenFF protein ligand benchmark dataset using perses. You must be located in the benchmarks subdirectory of the cloned perses repository to use this tool.

    To run a specific transformation (edge) for a specific protein (target), use the --target and --edge flags. The edge index is zero-based.

    If a calculation fails, re-running the exact same command will resume the calculation.

    # Run the seventh edge (index 6) for the tyk2 protein
python run_benchmarks.py --target tyk2 --edge 6

# Resume a failed edge
python run_benchmarks.py --target tyk2 --edge 6

  12. Install espaloma via mamba

    main

    It is recommended to use mamba for installation as it is a faster drop-in replacement for conda. You can create a dedicated environment named espaloma with the specific version 0.3.2 from the conda-forge channel using the following command:

    $ mamba create --name espaloma -c conda-forge "espaloma=0.3.2"