Merge pull request #2 from CompOmics/pretrained_model_choose

Black improvements

Author: Alirezak2n

ideeplc/__init__.py

@@ -1,3 +1,3 @@
 """iDeepLC: A deep Learning-based retention time predictor for unseen modified peptides with a novel encoding system"""
 
-__version__ = "1.2.1"
+__version__ = "1.3.0"

tests/test_CLI.py

@@ -2,6 +2,7 @@
 import unittest
 from ideeplc.__main__ import _argument_parser
 
+
 class TestCLI(unittest.TestCase):
     def test_argument_parser(self):
         """Test the CLI argument parser."""
@@ -13,5 +14,6 @@ def test_argument_parser(self):
         self.assertTrue(args.save, "Save argument should be True")
         self.assertTrue(args.calibrate, "Calibrate argument should be True")
 
+
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()

tests/test_config.py

@@ -1,5 +1,6 @@
 from ideeplc.config import get_config
 
+
 def test_get_config():
     """Test the get_config function to ensure it returns a valid configuration."""
     config = get_config(epoch=5)  # Adjust epoch for testing purposes
@@ -8,4 +9,5 @@ def test_get_config():
     assert config["epochs"] == 5, "Epoch in config should match the input value"
     assert "learning_rate" in config, "'learning_rate' should be in config"
 
-test_get_config()
+
+test_get_config()

tests/test_data_initialize.py

@@ -2,21 +2,26 @@
 import pandas as pd
 from ideeplc.data_initialize import data_initialize
 
+
 def test_data_initialize():
     """Test the data_initialize function."""
-    test_csv_path = "ideeplc/example_input/Hela_deeprt.csv"  # Path to a sample test CSV file
+    test_csv_path = (
+        "ideeplc/example_input/Hela_deeprt.csv"  # Path to a sample test CSV file
+    )
     matrix_input, x_shape = data_initialize(csv_path=test_csv_path)
 
     df_3 = pd.DataFrame(matrix_input[3][0])
-    expected = pd.DataFrame({
-        1: [1.286264, 0.209934, 12.0, 23.0, 2.0, 5.0],
-        2: [-1.076330, -2.715786, 6.0, 11.0, 7.0, 2.0],
-        3: [-1.435417, 1.267611, 14.0, 16.0, 4.0, 2.0],
-        4: [-1.280370, -2.870834, 6.0, 10.0, 4.0, 2.0]
-    }, index=[0, 1, 2, 3, 4, 5])
+    expected = pd.DataFrame(
+        {
+            1: [1.286264, 0.209934, 12.0, 23.0, 2.0, 5.0],
+            2: [-1.076330, -2.715786, 6.0, 11.0, 7.0, 2.0],
+            3: [-1.435417, 1.267611, 14.0, 16.0, 4.0, 2.0],
+            4: [-1.280370, -2.870834, 6.0, 10.0, 4.0, 2.0],
+        },
+        index=[0, 1, 2, 3, 4, 5],
+    )
 
     pd.testing.assert_frame_equal(df_3.iloc[:6, 1:5], expected, check_dtype=False)
     assert matrix_input is not None, "Matrix input should not be None"
     assert isinstance(x_shape, tuple), "x_shape should be a tuple"
     assert x_shape == (1, 41, 62), "Expected shape of x is (1, 41, 62)"
-

tests/test_model_save_path.py

@@ -8,6 +8,6 @@ def test_get_model_save_path():
     pretrained_model, model_dir = get_model_save_path()
     assert isinstance(pretrained_model, Path), "Model path should be a Path object"
     assert isinstance(model_dir, Path), "Model directory should be a Path object"
-    assert pretrained_model.name == "pretrained_model.pth", "Model name should be 'pretrained_model.pth'"
-
-
+    assert (
+        pretrained_model.name == "pretrained_model.pth"
+    ), "Model name should be 'pretrained_model.pth'"

tests/test_prediction_calibration.py

@@ -7,29 +7,35 @@
 from ideeplc.config import get_config
 from ideeplc.ideeplc_core import get_model_save_path
 
+
 def test_predict():
     """Test the predict function."""
     # Mock data and model
     config = get_config()
     pretrained_model, model_dir = get_model_save_path()
 
-    test_csv_path = "ideeplc/example_input/Hela_deeprt.csv"  # Path to a sample test CSV file
+    test_csv_path = (
+        "ideeplc/example_input/Hela_deeprt.csv"  # Path to a sample test CSV file
+    )
     matrix_input, x_shape = data_initialize(csv_path=test_csv_path)
-    dataloader = DataLoader(matrix_input, batch_size=config["batch_size"], shuffle=False)
+    dataloader = DataLoader(
+        matrix_input, batch_size=config["batch_size"], shuffle=False
+    )
     model = MyNet(x_shape=x_shape, config=config)
     device = torch.device("cpu")
-    model.load_state_dict(torch.load(pretrained_model, map_location=device), strict=False)
+    model.load_state_dict(
+        torch.load(pretrained_model, map_location=device), strict=False
+    )
     loss_fn = torch.nn.L1Loss()
 
-
     pred_loss, pred_cor, pred_results, ground_truth = predict(
         model=model,
         dataloader_input=dataloader,
         loss_fn=loss_fn,
         device=device,
         calibrate=False,
         input_file=test_csv_path,
-        save_results=False
+        save_results=False,
     )
     assert pred_loss is not None, "Prediction loss should not be None"
     assert pred_results is not None, "Prediction results should not be None"
@@ -43,9 +49,10 @@ def test_predict():
         device=device,
         calibrate=True,
         input_file=test_csv_path,
-        save_results=False
+        save_results=False,
     )
 
     assert pred_loss < 200, "Calibrated prediction loss should be less than 200"
-    assert pred_cor > 0.95, "Calibrated prediction correlation should be greater than 0.95"
-
+    assert (
+        pred_cor > 0.95
+    ), "Calibrated prediction correlation should be greater than 0.95"