chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,301 @@
|
||||
"""Torch Module for Atomic Convolution Layer"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import numpy as np
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class RadialPooling(nn.Module):
|
||||
r"""Radial pooling from `Atomic Convolutional Networks for
|
||||
Predicting Protein-Ligand Binding Affinity <https://arxiv.org/abs/1703.10603>`__
|
||||
|
||||
We denote the distance between atom :math:`i` and :math:`j` by :math:`r_{ij}`.
|
||||
|
||||
A radial pooling layer transforms distances with radial filters. For radial filter
|
||||
indexed by :math:`k`, it projects edge distances with
|
||||
|
||||
.. math::
|
||||
h_{ij}^{k} = \exp(-\gamma_{k}|r_{ij}-r_{k}|^2)
|
||||
|
||||
If :math:`r_{ij} < c_k`,
|
||||
|
||||
.. math::
|
||||
f_{ij}^{k} = 0.5 * \cos(\frac{\pi r_{ij}}{c_k} + 1),
|
||||
|
||||
else,
|
||||
|
||||
.. math::
|
||||
f_{ij}^{k} = 0.
|
||||
|
||||
Finally,
|
||||
|
||||
.. math::
|
||||
e_{ij}^{k} = h_{ij}^{k} * f_{ij}^{k}
|
||||
|
||||
Parameters
|
||||
----------
|
||||
interaction_cutoffs : float32 tensor of shape (K)
|
||||
:math:`c_k` in the equations above. Roughly they can be considered as learnable cutoffs
|
||||
and two atoms are considered as connected if the distance between them is smaller than
|
||||
the cutoffs. K for the number of radial filters.
|
||||
rbf_kernel_means : float32 tensor of shape (K)
|
||||
:math:`r_k` in the equations above. K for the number of radial filters.
|
||||
rbf_kernel_scaling : float32 tensor of shape (K)
|
||||
:math:`\gamma_k` in the equations above. K for the number of radial filters.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, interaction_cutoffs, rbf_kernel_means, rbf_kernel_scaling
|
||||
):
|
||||
super(RadialPooling, self).__init__()
|
||||
|
||||
self.interaction_cutoffs = nn.Parameter(
|
||||
interaction_cutoffs.reshape(-1, 1, 1), requires_grad=True
|
||||
)
|
||||
self.rbf_kernel_means = nn.Parameter(
|
||||
rbf_kernel_means.reshape(-1, 1, 1), requires_grad=True
|
||||
)
|
||||
self.rbf_kernel_scaling = nn.Parameter(
|
||||
rbf_kernel_scaling.reshape(-1, 1, 1), requires_grad=True
|
||||
)
|
||||
|
||||
def forward(self, distances):
|
||||
"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Apply the layer to transform edge distances.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
distances : Float32 tensor of shape (E, 1)
|
||||
Distance between end nodes of edges. E for the number of edges.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Float32 tensor of shape (K, E, 1)
|
||||
Transformed edge distances. K for the number of radial filters.
|
||||
"""
|
||||
scaled_euclidean_distance = (
|
||||
-self.rbf_kernel_scaling * (distances - self.rbf_kernel_means) ** 2
|
||||
) # (K, E, 1)
|
||||
rbf_kernel_results = th.exp(scaled_euclidean_distance) # (K, E, 1)
|
||||
|
||||
cos_values = 0.5 * (
|
||||
th.cos(np.pi * distances / self.interaction_cutoffs) + 1
|
||||
) # (K, E, 1)
|
||||
cutoff_values = th.where(
|
||||
distances <= self.interaction_cutoffs,
|
||||
cos_values,
|
||||
th.zeros_like(cos_values),
|
||||
) # (K, E, 1)
|
||||
|
||||
# Note that there appears to be an inconsistency between the paper and
|
||||
# DeepChem's implementation. In the paper, the scaled_euclidean_distance first
|
||||
# gets multiplied by cutoff_values, followed by exponentiation. Here we follow
|
||||
# the practice of DeepChem.
|
||||
return rbf_kernel_results * cutoff_values
|
||||
|
||||
|
||||
def msg_func(edges):
|
||||
"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Send messages along edges.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
edges : EdgeBatch
|
||||
A batch of edges.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict mapping 'm' to Float32 tensor of shape (E, K * T)
|
||||
Messages computed. E for the number of edges, K for the number of
|
||||
radial filters and T for the number of features to use
|
||||
(types of atomic number in the paper).
|
||||
"""
|
||||
return {
|
||||
"m": th.einsum("ij,ik->ijk", edges.src["hv"], edges.data["he"]).view(
|
||||
len(edges), -1
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
def reduce_func(nodes):
|
||||
"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Collect messages and update node representations.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nodes : NodeBatch
|
||||
A batch of nodes.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict mapping 'hv_new' to Float32 tensor of shape (V, K * T)
|
||||
Updated node representations. V for the number of nodes, K for the number of
|
||||
radial filters and T for the number of features to use
|
||||
(types of atomic number in the paper).
|
||||
"""
|
||||
return {"hv_new": nodes.mailbox["m"].sum(1)}
|
||||
|
||||
|
||||
class AtomicConv(nn.Module):
|
||||
r"""Atomic Convolution Layer from `Atomic Convolutional Networks for
|
||||
Predicting Protein-Ligand Binding Affinity <https://arxiv.org/abs/1703.10603>`__
|
||||
|
||||
Denoting the type of atom :math:`i` by :math:`z_i` and the distance between atom
|
||||
:math:`i` and :math:`j` by :math:`r_{ij}`.
|
||||
|
||||
**Distance Transformation**
|
||||
|
||||
An atomic convolution layer first transforms distances with radial filters and
|
||||
then perform a pooling operation.
|
||||
|
||||
For radial filter indexed by :math:`k`, it projects edge distances with
|
||||
|
||||
.. math::
|
||||
h_{ij}^{k} = \exp(-\gamma_{k}|r_{ij}-r_{k}|^2)
|
||||
|
||||
If :math:`r_{ij} < c_k`,
|
||||
|
||||
.. math::
|
||||
f_{ij}^{k} = 0.5 * \cos(\frac{\pi r_{ij}}{c_k} + 1),
|
||||
|
||||
else,
|
||||
|
||||
.. math::
|
||||
f_{ij}^{k} = 0.
|
||||
|
||||
Finally,
|
||||
|
||||
.. math::
|
||||
e_{ij}^{k} = h_{ij}^{k} * f_{ij}^{k}
|
||||
|
||||
**Aggregation**
|
||||
|
||||
For each type :math:`t`, each atom collects distance information from all neighbor atoms
|
||||
of type :math:`t`:
|
||||
|
||||
.. math::
|
||||
p_{i, t}^{k} = \sum_{j\in N(i)} e_{ij}^{k} * 1(z_j == t)
|
||||
|
||||
Then concatenate the results for all RBF kernels and atom types.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
interaction_cutoffs : float32 tensor of shape (K)
|
||||
:math:`c_k` in the equations above. Roughly they can be considered as learnable cutoffs
|
||||
and two atoms are considered as connected if the distance between them is smaller than
|
||||
the cutoffs. K for the number of radial filters.
|
||||
rbf_kernel_means : float32 tensor of shape (K)
|
||||
:math:`r_k` in the equations above. K for the number of radial filters.
|
||||
rbf_kernel_scaling : float32 tensor of shape (K)
|
||||
:math:`\gamma_k` in the equations above. K for the number of radial filters.
|
||||
features_to_use : None or float tensor of shape (T)
|
||||
In the original paper, these are atomic numbers to consider, representing the types
|
||||
of atoms. T for the number of types of atomic numbers. Default to None.
|
||||
|
||||
Note
|
||||
----
|
||||
|
||||
* This convolution operation is designed for molecular graphs in Chemistry, but it might
|
||||
be possible to extend it to more general graphs.
|
||||
|
||||
* There seems to be an inconsistency about the definition of :math:`e_{ij}^{k}` in the
|
||||
paper and the author's implementation. We follow the author's implementation. In the
|
||||
paper, :math:`e_{ij}^{k}` was defined as
|
||||
:math:`\exp(-\gamma_{k}|r_{ij}-r_{k}|^2 * f_{ij}^{k})`.
|
||||
|
||||
* :math:`\gamma_{k}`, :math:`r_k` and :math:`c_k` are all learnable.
|
||||
|
||||
Example
|
||||
-------
|
||||
>>> import dgl
|
||||
>>> import numpy as np
|
||||
>>> import torch as th
|
||||
>>> from dgl.nn import AtomicConv
|
||||
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> feat = th.ones(6, 1)
|
||||
>>> edist = th.ones(6, 1)
|
||||
>>> interaction_cutoffs = th.ones(3).float() * 2
|
||||
>>> rbf_kernel_means = th.ones(3).float()
|
||||
>>> rbf_kernel_scaling = th.ones(3).float()
|
||||
>>> conv = AtomicConv(interaction_cutoffs, rbf_kernel_means, rbf_kernel_scaling)
|
||||
>>> res = conv(g, feat, edist)
|
||||
>>> res
|
||||
tensor([[0.5000, 0.5000, 0.5000],
|
||||
[0.5000, 0.5000, 0.5000],
|
||||
[0.5000, 0.5000, 0.5000],
|
||||
[1.0000, 1.0000, 1.0000],
|
||||
[0.5000, 0.5000, 0.5000],
|
||||
[0.0000, 0.0000, 0.0000]], grad_fn=<ViewBackward>)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
interaction_cutoffs,
|
||||
rbf_kernel_means,
|
||||
rbf_kernel_scaling,
|
||||
features_to_use=None,
|
||||
):
|
||||
super(AtomicConv, self).__init__()
|
||||
|
||||
self.radial_pooling = RadialPooling(
|
||||
interaction_cutoffs=interaction_cutoffs,
|
||||
rbf_kernel_means=rbf_kernel_means,
|
||||
rbf_kernel_scaling=rbf_kernel_scaling,
|
||||
)
|
||||
if features_to_use is None:
|
||||
self.num_channels = 1
|
||||
self.features_to_use = None
|
||||
else:
|
||||
self.num_channels = len(features_to_use)
|
||||
self.features_to_use = nn.Parameter(
|
||||
features_to_use, requires_grad=False
|
||||
)
|
||||
|
||||
def forward(self, graph, feat, distances):
|
||||
"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Apply the atomic convolution layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
Topology based on which message passing is performed.
|
||||
feat : Float32 tensor of shape :math:`(V, 1)`
|
||||
Initial node features, which are atomic numbers in the paper.
|
||||
:math:`V` for the number of nodes.
|
||||
distances : Float32 tensor of shape :math:`(E, 1)`
|
||||
Distance between end nodes of edges. E for the number of edges.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Float32 tensor of shape :math:`(V, K * T)`
|
||||
Updated node representations. :math:`V` for the number of nodes, :math:`K` for the
|
||||
number of radial filters, and :math:`T` for the number of types of atomic numbers.
|
||||
"""
|
||||
with graph.local_scope():
|
||||
radial_pooled_values = self.radial_pooling(distances).to(
|
||||
feat
|
||||
) # (K, E, 1)
|
||||
if self.features_to_use is not None:
|
||||
feat = (feat == self.features_to_use).to(feat) # (V, T)
|
||||
graph.ndata["hv"] = feat
|
||||
graph.edata["he"] = radial_pooled_values.transpose(1, 0).squeeze(
|
||||
-1
|
||||
) # (E, K)
|
||||
graph.update_all(msg_func, reduce_func)
|
||||
|
||||
return graph.ndata["hv_new"].view(
|
||||
graph.num_nodes(), -1
|
||||
) # (V, K * T)
|
||||
Reference in New Issue
Block a user