Merge pull request #150 from Genentech/update-ledidi

fixed ledidi code to work with latest version and reran design tutorial

Author: avantikalal

docs/tutorials/4_design.ipynb

@@ -69,7 +69,7 @@ install_requires =
     captum == 0.5.0
     logomaker >= 0.8
     pyBigWig
-    ledidi < 2.0.0
+    ledidi
     tangermeme >= 0.5.0
     memelite
     pygenomeviz <= 0.4.4

setup.cfg

@@ -247,9 +247,6 @@ def ledidi(
     """
     from ledidi import Ledidi
 
-    # Add the prediction transform
-    model.add_transform(prediction_transform)
-
     def loss_func(x, target):
         return -Tensor(x).mean()
 
@@ -264,30 +261,40 @@ def loss_func(x, target):
     else:
         input_mask = None
 
+    # Add the prediction transform
+    model.add_transform(prediction_transform)
+
     # Move model to device
     orig_device = model.device
     model = model.to(torch.device(devices))
+    model.eval()
+
+    try:
+        print("Running Ledidi")
+
+        # Initialize ledidi
+        designer = Ledidi(
+            model,
+            X[0].shape,
+            output_loss=loss_func,
+            max_iter=max_iter,
+            input_mask=input_mask,
+            target=None,
+            **kwargs,
+        )
+        designer = designer.to(torch.device(devices))
 
-    # Initialize ledidi
-    designer = Ledidi(
-        model,
-        X[0].shape,
-        output_loss=loss_func,
-        max_iter=max_iter,
-        input_mask=input_mask,
-        target=None,
-        **kwargs,
-    )
-    designer = designer.to(torch.device(devices))
+        # Run ledidi
+        X_hat = designer.fit_transform(X, torch.tensor(0)).cpu()
 
-    # Run ledidi
-    X_hat = designer.fit_transform(X, None).cpu()
+    finally:
+        print("Cleaning up model state...")
 
-    # Transfer device
-    model = model.to(orig_device)
+        # Transfer device
+        model = model.to(orig_device)
 
-    # Remove the transform
-    model.reset_transform()
+        # Remove the transform
+        model.reset_transform()
 
     # Return sequences as strings
     return convert_input_type(X_hat, "strings")

src/grelu/design.py

@@ -2,10 +2,11 @@
 import pandas as pd
 from torch import Tensor, nn
 
-from grelu.design import evolve
+from grelu.design import evolve, ledidi
 from grelu.lightning import LightningModel
 from grelu.transforms.prediction_transforms import Aggregate, Specificity
 from grelu.transforms.seq_transforms import PatternScore
+from grelu.sequence.format import check_string_dna
 
 model = LightningModel(
     model_params={
@@ -220,3 +221,19 @@ def test_evolve_7():
     assert np.all(output["iter"] == [0, 1, 2])
     assert np.all(output.seq == ['AT', 'AC', 'CC'])
     assert np.all(output.label1 == [0.5, 1.5, 2.])
+
+
+def test_ledidi():
+    output = ledidi(
+        seq='GGTATTCATT',
+        model= model,
+        prediction_transform = None,
+        max_iter = 20000,
+        positions = None,
+        devices = "cpu",
+        num_workers = 1,
+        lr=0.01,
+        early_stopping_iter=1000
+    )
+    assert isinstance(output, list)
+    assert check_string_dna(output)