Repository: AstraZeneca/chemicalx (language: python)

Issue:
Add the SSI-DDI Model

- Please read the paper first. It is [here](https://academic.oup.com/bib/article-abstract/22/6/bbab133/6265181?redirectedFrom=fulltext).
- There is also code-release with the paper [here](https://github.com/kanz76/SSI-DDI).
- After that read the contributing guidelines.
-   If there is an existing open-source version of the model please take a look.
-   ChemicalX is built on top of [PyTorch 1.10.](https://pytorch.org/) and [torchdrug](https://torchdrug.ai/).
- A similar model is `EPGCNDS` which uses`GraphConvolutions`  to generate drug representations. Take a look at the layer definition [here](https://github.com/AstraZeneca/chemicalx/blob/main/chemicalx/models/epgcnds.py). You should use the layers from `torchdrug` not the models.
- The library heavily builds on top on [torchdrug](https://torchdrug.ai/) and molecules in batches are PackedGraphs.
- There is already a model class under `./chemicalx/models/`
- Context features, drug level features, and labels are all FloatTensors.
- Look at the examples and tests under `./examples/` and `./tests/`.
- Add auxiliary layers as you see fit - please document these, add tests, and add these layers to the main readme.md if needed.
- Add typing to the initialization and forward pass.
- Non-data-dependent hyper-ammeters should have default values. 
- Please add tests under `./tests/` and make sure that your model/layer is tested with real data.
- Write an example under `./examples/`. What is the AUC on the test set? Is it reasonable?




Discussion / hints:
Where can I find the contribution guideline and looks like the input of the original model in that paper only has two graph data and different interaction types without context features, do I need to incorporate also the context features into the current framework? Or do I need to find a way to transform the context features into a specific interaction type? Sorry that this is my first time using such a cooperative mode of GitHub and there might be some silly questions.
1. Context features are unused in this model.
2. There is a public reference implementation out.
> 1. Context features are unused in this model.
> 2. There is a public reference implementation out.

Thanks! Where can I find the public reference implementation out? Is it the documentation of this ChemicalX or just the original code associated with this paper?
https://github.com/kanz76/SSI-DDI/blob/master/models.py

Do not rely on PyTorch Geometric.
Sure, I will try to build myself the Graph Attention Layer, which is one of the most important building block for SSI-DDI

On Feb 2, 2022, at 12:22 PM, Benedek Rozemberczki ***@***.******@***.***>> wrote:



https://github.com/kanz76/SSI-DDI/blob/master/models.py<https://nam04.safelinks.protection.outlook.com/?url=https%3A%2F%2Fgithub.com%2Fkanz76%2FSSI-DDI%2Fblob%2Fmaster%2Fmodels.py&data=04%7C01%7Cyu.wang.1%40Vanderbilt.Edu%7Cb51098f241e14c071cba08d9e6790151%7Cba5a7f39e3be4ab3b45067fa80faecad%7C0%7C0%7C637794229648404353%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000&sdata=EgYoJG2aVysTsNjx70olhWWrjJEEhEcCU2elXCjW93I%3D&reserved=0>

Do not rely on PyTorch Geometric.

—
Reply to this email directly, view it on GitHub<https://nam04.safelinks.protection.outlook.com/?url=https%3A%2F%2Fgithub.com%2FAstraZeneca%2Fchemicalx%2Fissues%2F11%23issuecomment-1028226190&data=04%7C01%7Cyu.wang.1%40Vanderbilt.Edu%7Cb51098f241e14c071cba08d9e6790151%7Cba5a7f39e3be4ab3b45067fa80faecad%7C0%7C0%7C637794229648404353%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000&sdata=l4aKvc6N%2BEQcYmZLGlfN8Pkd06SgAfAFbsIeqd5qCxQ%3D&reserved=0>, or unsubscribe<https://nam04.safelinks.protection.outlook.com/?url=https%3A%2F%2Fgithub.com%2Fnotifications%2Funsubscribe-auth%2FAM2OPGTDBC3QD6FZ6R6AEA3UZFY7HANCNFSM5LWMCK3A&data=04%7C01%7Cyu.wang.1%40Vanderbilt.Edu%7Cb51098f241e14c071cba08d9e6790151%7Cba5a7f39e3be4ab3b45067fa80faecad%7C0%7C0%7C637794229648404353%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000&sdata=cGVinCCZ1bWUYhoSLd8Lv2Rl6xTbb5z7emK9WS1GOHM%3D&reserved=0>.
Triage notifications on the go with GitHub Mobile for iOS<https://nam04.safelinks.protection.outlook.com/?url=https%3A%2F%2Fapps.apple.com%2Fapp%2Fapple-store%2Fid1477376905%3Fct%3Dnotification-email%26mt%3D8%26pt%3D524675&data=04%7C01%7Cyu.wang.1%40Vanderbilt.Edu%7Cb51098f241e14c071cba08d9e6790151%7Cba5a7f39e3be4ab3b45067fa80faecad%7C0%7C0%7C637794229648560567%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000&sdata=TxAOv%2BXCkUZENlgTOT3L6IDSHX%2BWke2yO1M4OC5CFNw%3D&reserved=0> or Android<https://nam04.safelinks.protection.outlook.com/?url=https%3A%2F%2Fplay.google.com%2Fstore%2Fapps%2Fdetails%3Fid%3Dcom.github.android%26referrer%3Dutm_campaign%253Dnotification-email%2526utm_medium%253Demail%2526utm_source%253Dgithub&data=04%7C01%7Cyu.wang.1%40Vanderbilt.Edu%7Cb51098f241e14c071cba08d9e6790151%7Cba5a7f39e3be4ab3b45067fa80faecad%7C0%7C0%7C637794229648560567%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000&sdata=oRJMMBAcD47UTfhmOaCAKyuJW%2BzhpJFIw2fZM%2FYgino%3D&reserved=0>.
You are receiving this because you commented.Message ID: ***@***.***>


Torchdrug already has that layer.
The original paper defines different M learnable matrices for different interaction types in Equation (11), I guess here in our problem, can we assume only one interaction type is considered since the input is only the leftmolecules and rightmolecules.
Yes.
Just finished coding the model part, which is all based on torchdrug, pytorch and the current existing packages of chemicals, and currently working on testing the model. One question is that does the current pipeline support GPU training? If not, I guess I will just modify it a little bit on my end and try to train using GPU to test my implementation.
For testing whether our implementation is right, can I directly write a testing case following the previous script written in the test folder on github? But one thing is since it currently does not support GPU training, it will take me so long time to get the results. Can someone help me with that?

Fix patch:
diff --git a/chemicalx/models/ssiddi.py b/chemicalx/models/ssiddi.py
index dd62c16..e4b4724 100644
--- a/chemicalx/models/ssiddi.py
+++ b/chemicalx/models/ssiddi.py
@@ -1,13 +1,116 @@
 """An implementation of the SSI-DDI model."""
 
-from .base import UnimplementedModel
+import torch
+import torch.nn.functional
+from torch.nn import LayerNorm
+from torch.nn.modules.container import ModuleList
+from torchdrug.data import PackedGraph
+from torchdrug.layers import GraphAttentionConv, MeanReadout
+
+from chemicalx.constants import TORCHDRUG_NODE_FEATURES
+from chemicalx.data import DrugPairBatch
+from chemicalx.models import Model
 
 __all__ = [
     "SSIDDI",
 ]
 
 
-class SSIDDI(UnimplementedModel):
+class EmbeddingLayer(torch.nn.Module):
+    """Attention layer."""
+
+    def __init__(self, feature_number: int):
+        """Initialize the relational embedding layer.
+
+        :param feature_number: Number of features.
+        """
+        super().__init__()
+        self.weights = torch.nn.Parameter(torch.zeros(feature_number, feature_number))
+        torch.nn.init.xavier_uniform_(self.weights)
+
+    def forward(
+        self,
+        left_representations: torch.FloatTensor,
+        right_representations: torch.FloatTensor,
+        alpha_scores: torch.FloatTensor,
+    ):
+        """
+        Make a forward pass with the drug representations.
+
+        :param left_representations: Left side drug representations.
+        :param right_representations: Right side drug representations.
+        :param alpha_scores: Attention scores.
+        :returns: Positive label scores vector.
+        """
+        attention = torch.nn.functional.normalize(self.weights, dim=-1)
+        left_representations = torch.nn.functional.normalize(left_representations, dim=-1)
+        right_representations = torch.nn.functional.normalize(right_representations, dim=-1)
+        attention = attention.view(-1, self.weights.shape[0], self.weights.shape[1])
+        scores = alpha_scores * (left_representations @ attention @ right_representations.transpose(-2, -1))
+        scores = scores.sum(dim=(-2, -1)).view(-1, 1)
+        return scores
+
+
+class DrugDrugAttentionLayer(torch.nn.Module):
+    """Co-attention layer for drug pairs."""
+
+    def __init__(self, feature_number: int):
+        """Initialize the co-attention layer.
+
+        :param feature_number: Number of input features.
+        """
+        super().__init__()
+        self.weight_query = torch.nn.Parameter(torch.zeros(feature_number, feature_number // 2))
+        self.weight_key = torch.nn.Parameter(torch.zeros(feature_number, feature_number // 2))
+        self.bias = torch.nn.Parameter(torch.zeros(feature_number // 2))
+        self.attention = torch.nn.Parameter(torch.zeros(feature_number // 2))
+        self.tanh = torch.nn.Tanh()
+
+        torch.nn.init.xavier_uniform_(self.weight_query)
+        torch.nn.init.xavier_uniform_(self.weight_key)
+        torch.nn.init.xavier_uniform_(self.bias.view(*self.bias.shape, -1))
+        torch.nn.init.xavier_uniform_(self.attention.view(*self.attention.shape, -1))
+
+    def forward(self, left_representations: torch.Tensor, right_representations: torch.Tensor):
+        """Make a forward pass with the co-attention calculation.
+
+        :param left_representations: Matrix of left hand side representations.
+        :param right_representations: Matrix of right hand side representations.
+        :returns: Attention scores.
+        """
+        keys = left_representations @ self.weight_key
+        queries = right_representations @ self.weight_query
+        e_activations = queries.unsqueeze(-3) + keys.unsqueeze(-2) + self.bias
+        attentions = self.tanh(e_activations) @ self.attention
+        return attentions
+
+
+class SSIDDIBlock(torch.nn.Module):
+    """SSIDDI Block with convolution and pooling."""
+
+    def __init__(self, head_number: int, in_channels: int, out_channels: int):
+        """Initialize an SSI-DDI Block.
+
+        :param head_number: Number of attention heads.
+        :param in_channels: Number of input channels.
+        :param out_channels: Number of convolutional filters.
+        """
+        super().__init__()
+        self.conv = GraphAttentionConv(input_dim=in_channels, output_dim=out_channels, num_head=head_number)
+        self.readout = MeanReadout()
+
+    def forward(self, molecules: PackedGraph):
+        """Make a forward pass.
+
+        :param molecules: A batch of graphs.
+        :returns: The molecules with updated atom states and the pooled representations.
+        """
+        molecules.node_feature = self.conv(molecules, molecules.node_feature)
+        h_graphs = self.readout(molecules, molecules.node_feature)
+        return molecules, h_graphs
+
+
+class SSIDDI(Model):
     """An implementation of the SSI-DDI model from [nyamabo2021]_.
 
     .. seealso:: This model was suggested in https://github.com/AstraZeneca/chemicalx/issues/11
@@ -16,3 +119,59 @@ class SSIDDI(UnimplementedModel):
        for drug–drug interaction prediction <https://doi.org/10.1093/bib/bbab133>`_.
        *Briefings in Bioinformatics*, 22(6).
     """
+
+    def __init__(
+        self,
+        *,
+        molecule_channels: int = TORCHDRUG_NODE_FEATURES,
+        hidden_channels=(32, 32),
+        head_number=(2, 2),
+    ):
+        """Instantiate the SSI-DDI model.
+
+        :param molecule_channels: The number of molecular features.
+        :param hidden_channels: The list of neurons for each hidden layer block.
+        :param head_number: The number of attention heads in each block.
+        """
+        super().__init__()
+        self.initial_norm = LayerNorm(molecule_channels)
+        self.blocks = ModuleList()
+        self.net_norms = ModuleList()
+
+        channels = molecule_channels
+        for hidden_channel, head_number in zip(hidden_channels, head_number):
+            self.blocks.append(SSIDDIBlock(head_number, channels, hidden_channel))
+            self.net_norms.append(LayerNorm(hidden_channel))
+            channels = hidden_channel
+
+        self.co_attention = DrugDrugAttentionLayer(channels)
+        self.relational_embedding = EmbeddingLayer(channels)
+
+    def unpack(self, batch: DrugPairBatch):
+        """Return the left molecular graph and right molecular graph."""
+        return (
+            batch.drug_molecules_left,
+            batch.drug_molecules_right,
+        )
+
+    def _forward_molecules(self, molecules: PackedGraph):
+        molecules.node_feature = self.initial_norm(molecules.node_feature)
+        representation = []
+        for block, net_norm in zip(self.blocks, self.net_norms):
+            molecules, pooled_hidden_left = block(molecules)
+            representation.append(pooled_hidden_left)
+            molecules.node_feature = torch.nn.functional.elu(net_norm(molecules.node_feature))
+        return torch.stack(representation, dim=-2)
+
+    def forward(self, molecules_left: PackedGraph, molecules_right: PackedGraph) -> torch.FloatTensor:
+        """Run a forward pass of the SSI-DDI model.
+
+        :param molecules_left: Batched molecules for the left side drugs.
+        :param molecules_right: Batched molecules for the right side drugs.
+        :returns: A column vector of predicted synergy scores.
+        """
+        features_left = self._forward_molecules(molecules_left)
+        features_right = self._forward_molecules(molecules_right)
+        attentions = self.co_attention(features_left, features_right)
+        scores = torch.sigmoid(self.relational_embedding(features_left, features_right, attentions))
+        return scores
diff --git a/examples/ssiddi_examples.py b/examples/ssiddi_examples.py
new file mode 100644
index 0000000..d681a69
--- /dev/null
+++ b/examples/ssiddi_examples.py
@@ -0,0 +1,26 @@
+"""Example with SSI-DDI."""
+
+from chemicalx import pipeline
+from chemicalx.data import DrugCombDB
+from chemicalx.models import SSIDDI
+
+
+def main():
+    """Train and evaluate the SSIDDI model."""
+    dataset = DrugCombDB()
+    model = SSIDDI()
+    results = pipeline(
+        dataset=dataset,
+        model=model,
+        optimizer_kwargs=dict(lr=0.01, weight_decay=10**-7),
+        batch_size=1024,
+        epochs=20,
+        context_features=True,
+        drug_features=True,
+        drug_molecules=True,
+    )
+    results.summarize()
+
+
+if __name__ == "__main__":
+    main()


Test patch:
diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py
index 5c1c2a8..5904e5f 100644
--- a/tests/unit/test_models.py
+++ b/tests/unit/test_models.py
@@ -208,8 +208,18 @@ def test_mrgnn(self):
 
     def test_ssiddi(self):
         """Test SSIDDI."""
-        model = SSIDDI(x=2)
-        assert model.x == 2
+        model = SSIDDI()
+
+        optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=0.0001)
+        model.train()
+        loss = torch.nn.BCELoss()
+        for batch in self.generator:
+            optimizer.zero_grad()
+            prediction = model(batch.drug_molecules_left, batch.drug_molecules_right)
+            output = loss(prediction, batch.labels)
+            output.backward()
+            optimizer.step()
+            assert prediction.shape[0] == batch.labels.shape[0]
 
     def test_deepddi(self):
         """Test DeepDDI."""