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.

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

Class ii #74

Open
wants to merge 6 commits into
base: janossy
Choose a base branch
from
Open
Changes from all commits
Commits
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
2 changes: 1 addition & 1 deletion espaloma/mm/angle.py
Original file line number Diff line number Diff line change
@@ -30,8 +30,8 @@ def harmonic_angle(x, k, eq):

# NOTE:
# 0.25 because all angles are calculated twice
return 0.5 * esp.mm.functional.harmonic(x=x, k=k, eq=eq)

return 0.5 * esp.mm.functional.harmonic(x=x, k=k, eq=eq)

def linear_mixture_angle(x, coefficients, phases):
""" Angle energy with Linear basis function.
64 changes: 63 additions & 1 deletion espaloma/mm/energy.py
Original file line number Diff line number Diff line change
@@ -40,10 +40,34 @@ def apply_angle(nodes, suffix=""):
x=nodes.data["x"],
k=nodes.data["k%s" % suffix],
eq=nodes.data["eq%s" % suffix],
),
"u_urey_bradley%s"
% suffix: esp.mm.angle.urey_bradley(
x_between=nodes.data["x_between"],
k_urey_bradley=nodes.data["k_urey_bradley"],
eq_urey_bradley=nodes.data["eq_urey_bradley"],
),
"u_bond_bond%s"
% suffix: esp.mm.angle.bond_bond(
x_left=nodes.data["x_left"],
x_right=nodes.data["x_right"],
eq_left=nodes.data["eq_left"],
eq_right=nodes.data["eq_right"],
k_bond_bond=nodes.data["k_bond_bond"],
),
"u_bond_angle%s"
% suffix: esp.mm.angle.bond_angle(
x_left=nodes.data["x_left"],
x_right=nodes.data["x_right"],
x_angle=nodes.data["x"],
eq_left=nodes.data["eq_left"],
eq_right=nodes.data["eq_right"],
eq_angle=nodes.data["eq"],
k_bond_angle_left=nodes.data["k_bond_angle_left"],
k_bond_angle_right=nodes.data["k_bond_angle_right"],
)
}


def apply_angle_ii(nodes, suffix=""):
return {
"u_urey_bradley%s"
@@ -76,6 +100,23 @@ def apply_angle_ii(nodes, suffix=""):
def apply_torsion_ii(nodes, suffix=""):
""" Torsion energy in nodes. """
return {
"u_angle_angle%s"
% suffix: esp.mm.torsion.angle_angle(
x_angle_left=nodes.data["x_angle_left"],
x_angle_right=nodes.data["x_angle_right"],
eq_angle_left=nodes.data["eq_angle_left"],
eq_angle_right=nodes.data["eq_angle_right"],
k_angle_angle=nodes.data["k_angle_angle"],
),
"u_angle_angle_torsion%s"
% suffix: esp.mm.torsion.angle_angle_torsion(
x_angle_left=nodes.data["x_angle_left"],
x_angle_right=nodes.data["x_angle_right"],
eq_angle_left=nodes.data["eq_angle_left"],
eq_angle_right=nodes.data["eq_angle_right"],
x=nodes.data["x"],
k=nodes.data["k%s" % suffix],
),
"u_angle_angle%s"
% suffix: esp.mm.torsion.angle_angle(
x_angle_left=nodes.data["x_angle_left"],
@@ -316,6 +357,27 @@ def energy_in_graph(
ntype="onefour",
)

if class_ii is True:
g.apply_nodes(
lambda node: {
'u%s' % suffix:
node.data['u_urey_bradley%s' % suffix]\
+ node.data['u_bond_bond%s' % suffix]\
+ node.data['u_bond_angle%s' % suffix]
},
ntype='n3'
)

g.apply_nodes(
lambda node: {
'u%s' % suffix:
node.data['u_angle_angle%s' % suffix]\
+ node.data['u_angle_angle_torsion%s' % suffix]\
+ node.data['u_bond_torsion%s' % suffix]
},
ntype='n4'
)

# sum up energy
# bonded
g.multi_update_all(
1 change: 0 additions & 1 deletion espaloma/mm/functional.py
Original file line number Diff line number Diff line change
@@ -285,7 +285,6 @@ def linear_mixture(x, coefficients, phases=[0.0, 1.0]):

return u


def harmonic_periodic_coupled(
x_harmonic,
x_periodic,
19 changes: 11 additions & 8 deletions espaloma/mm/geometry.py
Original file line number Diff line number Diff line change
@@ -30,6 +30,9 @@ def _copy_src(edges, src=src, out=out):
# =============================================================================
def distance(x0, x1):
""" Distance. """
# TODO:
# assertions / shape docstrings
# p=2 indicates Euclidian distance
return torch.norm(x0 - x1, p=2, dim=-1)


@@ -55,7 +58,6 @@ def _dihedral(r0, r1):
""" Dihedral between normal vectors. """
return _angle(r0, r1)


def dihedral(
x0: torch.Tensor, x1: torch.Tensor, x2: torch.Tensor, x3: torch.Tensor
) -> torch.Tensor:
@@ -90,14 +92,19 @@ def dihedral(
return theta


return torch.atan2(top, bottom)

# =============================================================================
# GEOMETRY IN HYPERNODES
# =============================================================================
def apply_bond(nodes):
""" Bond length in nodes. """

return {"x": distance(x0=nodes.data["xyz0"], x1=nodes.data["xyz1"])}

return {
"x": distance(
x0=nodes.data["xyz0"],
x1=nodes.data["xyz1"]
),
}

def apply_angle(nodes):
""" Angle values in nodes. """
@@ -121,7 +128,6 @@ def apply_angle(nodes):
)
}


def apply_torsion(nodes):
""" Torsion dihedrals in nodes. """
return {
@@ -155,15 +161,12 @@ def apply_torsion(nodes):
)
}



# =============================================================================
# GEOMETRY IN GRAPH
# =============================================================================
# NOTE:
# The following functions modify graphs in-place.


def geometry_in_graph(g):
""" Assign values to geometric entities in graphs.

1 change: 0 additions & 1 deletion espaloma/mm/tests/test_energy.py
Original file line number Diff line number Diff line change
@@ -50,7 +50,6 @@ def test_energy():

esp.mm.energy.energy_in_graph(g, terms=["n2", "n3", "n4", "n4_improper"])


# def test_energy_consistent():
# g = esp.Graph("c1ccccc1")
#
87 changes: 87 additions & 0 deletions espaloma/mm/torsion.py
Original file line number Diff line number Diff line change
@@ -7,11 +7,13 @@
# =============================================================================
# MODULE FUNCTIONS
# =============================================================================

def periodic_torsion(
x, k, periodicity=list(range(1, 7)), phases=[0.0 for _ in range(6)]
):
""" Periodic torsion potential


Parameters
----------
x : `torch.Tensor`, `shape = (batch_size, 1)`
@@ -123,3 +125,88 @@ def bond_torsion(
k=k_right_torsion,
eq=eq_right_torsion,
)

def angle_angle(
x_angle_left,
x_angle_right,
eq_angle_left,
eq_angle_right,
k_angle_angle,
):
return esp.mm.functional.harmonic_harmonic_coupled(
x0=x_angle_left,
eq0=eq_angle_left,
x1=x_angle_right,
eq1=eq_angle_right,
k=k_angle_angle,
)

def angle_torsion(
x_angle_left,
x_angle_right,
eq_angle_left,
eq_angle_right,
x,
k_angle_torsion_left,
k_angle_torsion_right,
periodicity=list(range(1,3)),
):
return esp.mm.functional.harmonic_periodic_coupled(
x_harmonic=x_angle_left,
x_periodic=x,
k=k_angle_torsion_left,
eq=eq_angle_left,
periodicity=periodicity,
) + esp.mm.functional.harmonic_periodic_coupled(
x_harmonic=x_angle_right,
x_periodic=x,
k=k_angle_torsion_right,
eq=eq_angle_right,
periodicity=periodicity,
)

def angle_angle_torsion(
x_angle_left,
x_angle_right,
eq_angle_left,
eq_angle_right,
x,
k_angle_angle_torsion,
):
return esp.mm.functional.harmonic_harmonic_periodic_coupled(
theta0=x_angle_left,
theta1=x_angle_right,
eq0=eq_angle_left,
eq1=eq_angle_right,
phi=x,
k=k_angle_angle_torsion,
)

def bond_torsion(
x_bond_left,
x_bond_center,
x_bond_right,
x,
k_left_torsion,
k_center_torsion,
k_right_torsion,
eq_left_torsion,
eq_center_torsion,
eq_right_torsion,
):
return esp.mm.functional.harmonic_periodic_coupled(
x_harmonic=x_bond_left,
x_periodic=x,
k=k_left_torsion,
eq=eq_left_torsion,
) + esp.mm.functional.harmonic_periodic_coupled(
x_harmonic=x_bond_center,
x_periodic=x,
k=k_center_torsion,
eq=eq_center_torsion,
) + esp.mm.functional.harmonic_periodic_coupled(
x_harmonic=x_bond_right,
x_periodic=x,
k=k_right_torsion,
eq=eq_right_torsion,
)
Binary file added scripts/.DS_Store
Binary file not shown.
Loading