Repository: AstraZeneca/chemicalx (language: python)

Issue:
Add the DeepDrug Model
Dear @kajocina, 

- Please read the paper first. It is [here](https://www.biorxiv.org/content/biorxiv/early/2020/11/10/2020.11.09.375626.full.pdf).
- There is also code-release with the paper [here](https://github.com/wanwenzeng/deepdrug).
- 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:

Fix patch:
diff --git a/chemicalx/models/deepdrug.py b/chemicalx/models/deepdrug.py
index bf72ff6..0693785 100644
--- a/chemicalx/models/deepdrug.py
+++ b/chemicalx/models/deepdrug.py
@@ -1,13 +1,20 @@
 """An implementation of the DeepDrug model."""
 
-from .base import UnimplementedModel
+import torch
+from torch import nn
+from torchdrug.data import PackedGraph
+from torchdrug.layers import GraphConv, MaxReadout
+
+from chemicalx.constants import TORCHDRUG_NODE_FEATURES
+from chemicalx.data import DrugPairBatch
+from chemicalx.models import Model
 
 __all__ = [
     "DeepDrug",
 ]
 
 
-class DeepDrug(UnimplementedModel):
+class DeepDrug(Model):
     """An implementation of the DeepDrug model from [cao2020]_.
 
     .. seealso:: This model was suggested in https://github.com/AstraZeneca/chemicalx/issues/14
@@ -16,3 +23,71 @@ class DeepDrug(UnimplementedModel):
        for drug relation prediction <https://doi.org/10.1101/2020.11.09.375626>`_.
        *bioRxiv*, 2020.11.09.375626.
     """
+
+    def __init__(
+        self,
+        *,
+        molecule_channels: int = TORCHDRUG_NODE_FEATURES,
+        num_gcn_layers: int = 4,
+        gcn_layer_hidden_size: int = 64,
+        out_channels: int = 1,
+        dropout_rate: float = 0.1,
+    ):
+        """Instantiate the DeepDrug model.
+
+        :param molecule_channels: The number of molecular features.
+        :param num_gcn_layers: Number of GCN layers.
+        :param gcn_layer_hidden_size: number of hidden units in GCN layers
+        :param out_channels: The number of output channels.
+        :param dropout_rate: Dropout rate on the final fully-connected layer.
+        """
+        super(DeepDrug, self).__init__()
+        self.num_gcn_layers = num_gcn_layers
+        self.gcn_layer_hidden_size = gcn_layer_hidden_size
+        self.graph_convolution_first = GraphConv(molecule_channels, self.gcn_layer_hidden_size, batch_norm=True)
+
+        # add remaining GCN layers
+        self.layers = torch.nn.ModuleList(
+            GraphConv(self.gcn_layer_hidden_size, self.gcn_layer_hidden_size, batch_norm=True)
+            for _ in range(num_gcn_layers - 1)
+        )
+
+        self.readout = MaxReadout()
+        self.middle_channels = 2 * self.gcn_layer_hidden_size
+
+        self.final = nn.Sequential(
+            nn.BatchNorm1d(self.middle_channels),
+            nn.Dropout(p=dropout_rate),
+            nn.Linear(self.middle_channels, out_channels),
+            nn.Sigmoid(),
+        )
+
+    def unpack(self, batch: DrugPairBatch):
+        """Return the left drug molecules, and right drug molecules."""
+        return (
+            batch.drug_molecules_left,
+            batch.drug_molecules_right,
+        )
+
+    def _help_molecules(self, molecules: PackedGraph):
+        features = self.graph_convolution_first(molecules, molecules.data_dict["node_feature"])
+        for layer in self.layers:
+            features = layer(molecules, features)
+        features = self.readout(molecules, features)
+        return features
+
+    def forward(self, molecules_left: PackedGraph, molecules_right: PackedGraph) -> torch.FloatTensor:
+        """
+        Run a forward pass of the DeepDrug model.
+
+        :param molecules_left: Batched molecules for the left side drugs.
+        :param molecules_right: Batched molecules for the right side drugs.
+
+        :return: A column vector of predicted synergy scores.
+        """
+        features_left = self._help_molecules(molecules_left)
+        features_right = self._help_molecules(molecules_right)
+
+        # run the linear layer on concatenated left/right features
+        hidden = torch.cat([features_left, features_right], dim=1)
+        return self.final(hidden)
diff --git a/examples/deepdrug_examples.py b/examples/deepdrug_examples.py
new file mode 100644
index 0000000..ff82312
--- /dev/null
+++ b/examples/deepdrug_examples.py
@@ -0,0 +1,26 @@
+"""Example with DeepDrug."""
+
+from chemicalx import pipeline
+from chemicalx.data import DrugCombDB
+from chemicalx.models import DeepDrug
+
+
+def main():
+    """Train and evaluate the EPGCNDS model."""
+    dataset = DrugCombDB()
+    model = DeepDrug()
+    results = pipeline(
+        dataset=dataset,
+        model=model,
+        optimizer_kwargs=dict(lr=0.001),
+        batch_size=1024,
+        epochs=20,
+        context_features=False,
+        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 e863dd0..3d4bb5c 100644
--- a/tests/unit/test_models.py
+++ b/tests/unit/test_models.py
@@ -198,8 +198,18 @@ def test_deepddi(self):
 
     def test_deepdrug(self):
         """Test DeepDrug."""
-        model = DeepDrug(x=2)
-        assert model.x == 2
+        model = DeepDrug()
+
+        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+        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_deepsynergy(self):
         """Test DeepSynergy."""