Skip to content

Gsl setup #74

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

Open
wants to merge 10 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
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
13 changes: 13 additions & 0 deletions graph_structure_learning/graph_structure_learner.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
"""
Performs graph structure learning on the MSG as outlined in the
paper.
"""


class GraphStructureLearner:
"""
Implements GSL.
"""

def __init__(self):
pass
24 changes: 24 additions & 0 deletions graph_structure_learning/molecular_similarity_graph.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
"""
A class implementing the molecular similarity graph (MSG) by
taking data from the MoleculeProcessor class as well as the model
embeddings. Again, it should follow the steps in the notebook
and have methods that perform the following
- constructs an adjacency matrix from the tanimoto coefficients
subject to a cutoff,
- takes model embeddings and adjacency matrix and converts them
into PyTorch Geometric graph data.
- (optional) networkx visualization of molecular similarity graph
"""

class MolecularSimilarityGraph:
def __init__(self, moleculeData):
pass

def constructAdjacencyMatrix(self):
pass

def toGraphData(self):
pass

def visualize(self):
pass
8 changes: 8 additions & 0 deletions graph_structure_learning/run_with_gsl.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
"""
Run our GNN-GSL model.
- Much of this code can be taken from main.py, but
we will have to modify the running process so that it
takes the 'readout' output from the trainer and feeds
it into the GSL pipeline.
"""

26 changes: 26 additions & 0 deletions graph_structure_learning/xyz_processor.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
"""
Use Atomic Simulation Environment (ASE) to process the molecules
into individual .xyz files and then delete the files.
- Implement a 'MoleculeProcessor' object that takes a list of
JSON molecules on construction and has methods that perform the
exact same process as in our notebook. I.e., the following methods
must be implemented:
- converts the molecules into .xyz files,
- computes and stores all relevent metrics
(fingerprints, Tanimoto coeffs, etc.),
- deletes all .xyz files.
"""


class MoleculeProcessor:
def __init__(self, molList):
pass

def toXYZ(self):
pass

def computeMetrics(self):
pass

def teardownXYZ(self):
pass
144 changes: 82 additions & 62 deletions matdeeplearn/models/torchmd_et.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,9 @@
from matdeeplearn.models.torchmd_output_modules import Scalar, EquivariantScalar
from matdeeplearn.common.registry import registry
from matdeeplearn.preprocessor.helpers import node_rep_one_hot
@registry.register_model("torchmd_et")


@registry.register_model("torchmd_et")
class TorchMD_ET(BaseModel):
r"""The TorchMD equivariant Transformer architecture.

Expand Down Expand Up @@ -60,7 +60,7 @@ class TorchMD_ET(BaseModel):

def __init__(
self,
node_dim,
node_dim,
edge_dim,
output_dim,
hidden_channels=128,
Expand Down Expand Up @@ -110,7 +110,8 @@ def __init__(
self.distance_influence = distance_influence
self.max_z = max_z
self.pool = pool
assert pool_order in ['early', 'late'], f"{pool_order} is currently not supported"
assert pool_order in [
'early', 'late'], f"{pool_order} is currently not supported"
self.pool_order = pool_order
self.output_dim = output_dim
cutoff_lower = 0
Expand Down Expand Up @@ -159,10 +160,13 @@ def __init__(
self.post_lin_list = nn.ModuleList()
for i in range(self.num_post_layers):
if i == 0:
self.post_lin_list.append(nn.Linear(hidden_channels, post_hidden_channels))
self.post_lin_list.append(
nn.Linear(hidden_channels, post_hidden_channels))
else:
self.post_lin_list.append(nn.Linear(post_hidden_channels, post_hidden_channels))
self.post_lin_list.append(nn.Linear(post_hidden_channels, self.output_dim))
self.post_lin_list.append(
nn.Linear(post_hidden_channels, post_hidden_channels))
self.post_lin_list.append(
nn.Linear(post_hidden_channels, self.output_dim))

self.reset_parameters()

Expand All @@ -174,82 +178,93 @@ def reset_parameters(self):
for attn in self.attention_layers:
attn.reset_parameters()
self.out_norm.reset_parameters()

@conditional_grad(torch.enable_grad())
def _forward(self, data):

x = self.embedding(data.z)

#edge_index, edge_weight, edge_vec = self.distance(data.pos, data.batch)
#assert (
# edge_index, edge_weight, edge_vec = self.distance(data.pos, data.batch)
# assert (
# edge_vec is not None
#), "Distance module did not return directional information"
# ), "Distance module did not return directional information"
if self.otf_edge_index == True:
#data.edge_index, edge_weight, data.edge_vec, cell_offsets, offset_distance, neighbors = self.generate_graph(data, self.cutoff_radius, self.n_neighbors)
data.edge_index, data.edge_weight, data.edge_vec, _, _, _ = self.generate_graph(data, self.cutoff_radius, self.n_neighbors)
data.edge_attr = self.distance_expansion(data.edge_weight)

#mask = data.edge_index[0] != data.edge_index[1]
#data.edge_vec[mask] = data.edge_vec[mask] / torch.norm(data.edge_vec[mask], dim=1).unsqueeze(1)
data.edge_vec = data.edge_vec / torch.norm(data.edge_vec, dim=1).unsqueeze(1)

# data.edge_index, edge_weight, data.edge_vec, cell_offsets, offset_distance, neighbors = self.generate_graph(data, self.cutoff_radius, self.n_neighbors)
data.edge_index, data.edge_weight, data.edge_vec, _, _, _ = self.generate_graph(
data, self.cutoff_radius, self.n_neighbors)
data.edge_attr = self.distance_expansion(data.edge_weight)

# mask = data.edge_index[0] != data.edge_index[1]
# data.edge_vec[mask] = data.edge_vec[mask] / torch.norm(data.edge_vec[mask], dim=1).unsqueeze(1)
data.edge_vec = data.edge_vec / \
torch.norm(data.edge_vec, dim=1).unsqueeze(1)

if self.otf_node_attr == True:
data.x = node_rep_one_hot(data.z).float()
data.x = node_rep_one_hot(data.z).float()

if self.neighbor_embedding is not None:
x = self.neighbor_embedding(data.z, x, data.edge_index, data.edge_weight, data.edge_attr)
x = self.neighbor_embedding(
data.z, x, data.edge_index, data.edge_weight, data.edge_attr)

vec = torch.zeros(x.size(0), 3, x.size(1), device=x.device)

for attn in self.attention_layers:
dx, dvec = attn(x, vec, data.edge_index, data.edge_weight, data.edge_attr, data.edge_vec)
dx, dvec = attn(x, vec, data.edge_index,
data.edge_weight, data.edge_attr, data.edge_vec)
x = x + dx
vec = vec + dvec
# just output the embeddings => stop before the prediction layer
x = self.out_norm(x)

if self.prediction_level == "graph":
if self.pool_order == 'early':
x = getattr(torch_geometric.nn, self.pool)(x, data.batch)
for i in range(0, len(self.post_lin_list) - 1):
x = self.post_lin_list[i](x)
x = getattr(F, self.activation)(x)
x = self.post_lin_list[-1](x)
if self.pool_order == 'late':
x = getattr(torch_geometric.nn, self.pool)(x, data.batch)
#x = self.pool.pre_reduce(x, vec, data.z, data.pos, data.batch)
#x = self.pool.reduce(x, data.batch)
elif self.prediction_level == "node":
for i in range(0, len(self.post_lin_list) - 1):
x = self.post_lin_list[i](x)
x = getattr(F, self.activation)(x)
x = self.post_lin_list[-1](x)


# if self.prediction_level == "graph":
# if self.pool_order == 'early':
# x = getattr(torch_geometric.nn, self.pool)(x, data.batch)
# for i in range(0, len(self.post_lin_list) - 1):
# x = self.post_lin_list[i](x)
# x = getattr(F, self.activation)(x)
# x = self.post_lin_list[-1](x)
# if self.pool_order == 'late':
# x = getattr(torch_geometric.nn, self.pool)(x, data.batch)
# # x = self.pool.pre_reduce(x, vec, data.z, data.pos, data.batch)
# # x = self.pool.reduce(x, data.batch)
# elif self.prediction_level == "node":
# for i in range(0, len(self.post_lin_list) - 1):
# x = self.post_lin_list[i](x)
# x = getattr(F, self.activation)(x)
# x = self.post_lin_list[-1](x)

# TODO: FIGURE OUT HOW TO ACCESS EMBEDDINGS; WE NEED THEM TO COMPUTE
# MOLECULAR FINGERPRINTS.

return x

def forward(self, data):

output = {}
out = self._forward(data)
output["output"] = out
output["output"] = out

if self.gradient == True and out.requires_grad == True:
volume = torch.einsum("zi,zi->z", data.cell[:, 0, :], torch.cross(data.cell[:, 1, :], data.cell[:, 2, :], dim=1)).unsqueeze(-1)
# this is skipped reached since we're not getting the prediction (I think?)
# even if it is reached, we're probably fine lol.
if self.gradient == True and out.requires_grad == True:
volume = torch.einsum("zi,zi->z", data.cell[:, 0, :], torch.cross(
data.cell[:, 1, :], data.cell[:, 2, :], dim=1)).unsqueeze(-1)
grad = torch.autograd.grad(
out,
[data.pos, data.displacement],
grad_outputs=torch.ones_like(out),
create_graph=self.training)
out,
[data.pos, data.displacement],
grad_outputs=torch.ones_like(out),
create_graph=self.training)
forces = -1 * grad[0]
stress = grad[1]
stress = stress / volume.view(-1, 1, 1)
stress = stress / volume.view(-1, 1, 1)

output["pos_grad"] = forces
output["cell_grad"] = stress
output["pos_grad"] = forces
output["cell_grad"] = stress
else:
output["pos_grad"] = None
output["cell_grad"] = None
return output
output["pos_grad"] = None
output["cell_grad"] = None

return output

def __repr__(self):
return (
Expand All @@ -267,6 +282,7 @@ def __repr__(self):
f"cutoff_lower={self.cutoff_lower}, "
f"self.cutoff_radius={self.self.cutoff_radius})"
)

@property
def target_attr(self):
return "y"
Expand All @@ -285,7 +301,8 @@ def __init__(
cutoff_upper,
aggregation,
):
super(EquivariantMultiHeadAttention, self).__init__(aggr=aggregation, node_dim=0)
super(EquivariantMultiHeadAttention, self).__init__(
aggr=aggregation, node_dim=0)
assert hidden_channels % num_heads == 0, (
f"The number of hidden channels ({hidden_channels}) "
f"must be evenly divisible by the number of "
Expand All @@ -307,7 +324,8 @@ def __init__(
self.v_proj = nn.Linear(hidden_channels, hidden_channels * 3)
self.o_proj = nn.Linear(hidden_channels, hidden_channels * 3)

self.vec_proj = nn.Linear(hidden_channels, hidden_channels * 3, bias=False)
self.vec_proj = nn.Linear(
hidden_channels, hidden_channels * 3, bias=False)

self.dk_proj = None
if distance_influence in ["keys", "both"]:
Expand Down Expand Up @@ -343,21 +361,23 @@ def forward(self, x, vec, edge_index, r_ij, f_ij, d_ij):
k = self.k_proj(x).reshape(-1, self.num_heads, self.head_dim)
v = self.v_proj(x).reshape(-1, self.num_heads, self.head_dim * 3)

vec1, vec2, vec3 = torch.split(self.vec_proj(vec), self.hidden_channels, dim=-1)
vec1, vec2, vec3 = torch.split(
self.vec_proj(vec), self.hidden_channels, dim=-1)
vec = vec.reshape(-1, 3, self.num_heads, self.head_dim)
vec_dot = (vec1 * vec2).sum(dim=1)

dk = (
self.act(self.dk_proj(f_ij)).reshape(-1, self.num_heads, self.head_dim)
self.act(self.dk_proj(f_ij)).reshape(-1,
self.num_heads, self.head_dim)
if self.dk_proj is not None
else None
)
dv = (
self.act(self.dv_proj(f_ij)).reshape(-1, self.num_heads, self.head_dim * 3)
self.act(self.dv_proj(f_ij)).reshape(-1,
self.num_heads, self.head_dim * 3)
if self.dv_proj is not None
else None
)


# propagate_type: (q: Tensor, k: Tensor, v: Tensor, vec: Tensor, dk: Tensor, dv: Tensor, r_ij: Tensor, d_ij: Tensor)
x, vec = self.propagate(
Expand Down
Loading