Fix remaining issues from code review

RalfG · RalfG · commit 6e60ff2faa4a · 2026-03-18T11:13:38.000+01:00
diff --git a/im2deep/__main__.py b/im2deep/__main__.py
@@ -94,7 +94,7 @@ def cli(ctx, logging_level, profile, profile_name):
 
     console.print(_build_credits())
 
-
+# TODO:  Check that parameters match predict function in core
 # Implement psm_utils reading for calibration and prediction PSMLists
 @cli.command()
 @click.pass_context
diff --git a/im2deep/_io_helpers.py b/im2deep/_io_helpers.py
@@ -239,8 +239,7 @@ def validate_psm_list(psm_list: PSMList, needs_target: bool = False) -> PSMList:
     # Filter missing and high charge states (IM2Deep predictions are not reliable for charges >6)
     original_size = len(psm_list)
     charges = np.array([psm.peptidoform.precursor_charge for psm in psm_list])
-    psm_list_filtered = psm_list[charges != None]  # noqa: E711
-    psm_list_filtered = psm_list_filtered[charges <= 6]
+    psm_list_filtered = psm_list[(charges != None) & (charges <= 6)]  # noqa: E711
 
     # TODO: Is deepcopy really necessary or can it be avoided?
     psm_list_filtered = deepcopy(psm_list_filtered)
@@ -263,25 +262,20 @@ def validate_psm_list(psm_list: PSMList, needs_target: bool = False) -> PSMList:
             for psm in psm_list_filtered
         )
 
+        # TODO: Could be vectorized over all ion mobility values
         # If ion_mobility is present, convert to CCS
         for psm in psm_list_filtered:
-            if (
-                psm.ion_mobility is not None
-                and psm.metadata is not None
-                and psm.metadata.get("CCS") is None
-            ):
-                psm.metadata["CCS"] = str(
-                    im2ccs(
-                        psm.ion_mobility,
-                        psm.peptidoform.theoretical_mz,
-                        psm.peptidoform.precursor_charge,
+            if psm.ion_mobility is not None:
+                if psm.metadata is None:
+                    psm.metadata = {}
+                if "CCS" not in psm.metadata:
+                    psm.metadata["CCS"] = str(
+                        im2ccs(
+                            psm.ion_mobility,
+                            psm.peptidoform.theoretical_mz,
+                            psm.peptidoform.precursor_charge,
+                        )
                     )
-                )
-            # Ensure CCS is always stored as float
-            elif psm.metadata.get("CCS") is not None:
-                ccs_value = psm.metadata["CCS"]
-                if not isinstance(ccs_value, float):
-                    psm.metadata["CCS"] = float(ccs_value)
 
     if needs_target and not all_has_targets:
         raise IM2DeepError("PSMList must contain 'ion_mobility' or 'CCS' metadata for all PSMs.")
diff --git a/im2deep/_model_ops.py b/im2deep/_model_ops.py
@@ -130,6 +130,7 @@ def predict(
         config=_get_model_config(multi=multi),
         criterion=_get_loss_function(multi=multi),
     )
+    model.to(device)
     model.eval()
     LOGGER.debug(f"Model loaded on device: {device}")
 
diff --git a/im2deep/core.py b/im2deep/core.py
@@ -11,9 +11,9 @@
 from psm_utils.psm_list import PSMList
 
 from im2deep import _model_ops
+from im2deep._io_helpers import validate_psm_list
 from im2deep.calibration import Calibration, LinearCCSCalibration
 from im2deep.constants import DEFAULT_MODEL, DEFAULT_MULTI_MODEL
-from im2deep._io_helpers import validate_psm_list
 
 LOGGER = logging.getLogger(__name__)
 
diff --git a/tests/test_model_ops.py b/tests/test_model_ops.py
@@ -20,7 +20,8 @@ def test_load_model_from_path(self, temp_model_path):
         loaded_model = _model_ops.load_model(temp_model_path)
 
         assert isinstance(loaded_model, torch.nn.Module)
-        assert loaded_model.training is False  # Should be in eval mode by default
+        # load_model does not force eval mode; caller controls train/eval state.
+        assert loaded_model.training is model.training
 
     def test_load_model_from_module(self):
         """Test loading model from existing module."""
@@ -38,7 +39,7 @@ def test_load_model_none(self):
     def test_load_model_invalid_type(self):
         """Test loading model with invalid type raises TypeError."""
         with pytest.raises(TypeError):
-            _model_ops.load_model(12345)
+            _model_ops.load_model(12345)  # type: ignore
 
     def test_load_model_dict_checkpoint(self, temp_model_path):
         """Test loading model from dict checkpoint."""
@@ -169,7 +170,7 @@ def test_predict_loop_single_output(self):
 
         with patch("im2deep._model_ops.track", return_value=mock_data):
             predictions = _model_ops._predict_loop(
-                model=model, data_loader=mock_data, device="cpu"
+                model=model, data_loader=mock_data, device="cpu"  # type: ignore
             )
 
             assert isinstance(predictions, torch.Tensor)
@@ -190,7 +191,7 @@ def test_predict_loop_multi_output(self):
 
         with patch("im2deep._model_ops.track", return_value=mock_data):
             predictions = _model_ops._predict_loop(
-                model=model, data_loader=mock_data, device="cpu"
+                model=model, data_loader=mock_data, device="cpu"  # type: ignore
             )
 
             assert isinstance(predictions, torch.Tensor)
@@ -211,7 +212,7 @@ def test_predict_loop_no_grad(self):
         # Mock track to return our mock data
         with patch("im2deep._model_ops.track", return_value=mock_data):
             predictions = _model_ops._predict_loop(
-                model=model, data_loader=mock_data, device="cpu"
+                model=model, data_loader=mock_data, device="cpu"  # type: ignore
             )
 
             assert not predictions.requires_grad
@@ -220,14 +221,14 @@ def test_predict_loop_no_grad(self):
 class TestGetArchitecture:
     """Tests for _get_architecture function."""
 
-    @patch("im2deep._architecture.IM2Deep")
+    @patch("im2deep._model_ops.IM2Deep")
     def test_get_architecture_single(self, mock_im2deep):
         """Test getting single-output architecture."""
         arch = _model_ops._get_architecture(multi=False)
         # Should import IM2Deep
         assert arch is mock_im2deep
 
-    @patch("im2deep._architecture.IM2DeepMultiTransfer")
+    @patch("im2deep._model_ops.IM2DeepMultiTransfer")
     def test_get_architecture_multi(self, mock_multi):
         """Test getting multi-output architecture."""
         arch = _model_ops._get_architecture(multi=True)
@@ -257,7 +258,7 @@ def test_get_loss_function_single(self):
         loss = _model_ops._get_loss_function(multi=False)
         assert isinstance(loss, torch.nn.modules.loss._Loss)
 
-    @patch("im2deep._architecture.FlexibleLossSorted")
+    @patch("im2deep._model_ops.FlexibleLossSorted")
     def test_get_loss_function_multi(self, mock_loss):
         """Test getting multi-output loss function."""
         mock_instance = MagicMock()
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -5,13 +5,9 @@
 import pytest
 from psm_utils import PSMList
 
+from im2deep._io_helpers import parse_input, validate_psm_list
 from im2deep.exceptions import IM2DeepError
-from im2deep.utils import (
-    ccs2im,
-    im2ccs,
-    parse_input,
-    validate_psm_list,
-)
+from im2deep.utils import ccs2im, im2ccs
 
 
 class TestValidatePSMList:
@@ -28,6 +24,7 @@ def test_validate_psm_list_with_ccs(self, sample_psm_list_with_ccs):
         result = validate_psm_list(sample_psm_list_with_ccs, needs_target=True)
         assert isinstance(result, PSMList)
         for psm in result:
+            assert psm.metadata is not None
             assert "CCS" in psm.metadata
             # CCS should always be stored as float
             assert isinstance(psm.metadata["CCS"], float)
@@ -46,7 +43,7 @@ def test_validate_psm_list_empty(self):
     def test_validate_psm_list_not_psm_list(self):
         """Test validation fails with non-PSMList input."""
         with pytest.raises(IM2DeepError, match="PSMList"):
-            validate_psm_list([1, 2, 3])
+            validate_psm_list([1, 2, 3])  # type: ignore
 
 
 class TestParseInput:
@@ -131,6 +128,7 @@ def test_parse_input_legacy_format_detection(self, tmp_path):
         assert len(result) == 2
         # Check that CCS values are preserved in metadata
         for psm in result:
+            assert psm.metadata is not None
             assert "CCS" in psm.metadata
 
 
@@ -142,7 +140,7 @@ def test_ccs2im_basic(self):
         ccs = 450.0
         charge = 2
         mz = 500.0
-        im = ccs2im(ccs, charge, mz)
+        im = ccs2im(ccs, mz, charge)
 
         assert isinstance(im, float)
         assert im > 0
@@ -153,7 +151,7 @@ def test_ccs2im_array(self):
         charge = np.array([2, 3, 2])
         mz = np.array([500.0, 600.0, 550.0])
 
-        im = ccs2im(ccs, charge, mz)
+        im = ccs2im(ccs, mz, charge)
 
         assert isinstance(im, np.ndarray)
         assert len(im) == len(ccs)
@@ -164,7 +162,7 @@ def test_im2ccs_basic(self):
         im = 1.0
         charge = 2
         mz = 500.0
-        ccs = im2ccs(im, charge, mz)
+        ccs = im2ccs(im, mz, charge)
 
         assert isinstance(ccs, float)
         assert ccs > 0
@@ -175,7 +173,7 @@ def test_im2ccs_array(self):
         charge = np.array([2, 3, 2])
         mz = np.array([500.0, 600.0, 550.0])
 
-        ccs = im2ccs(im, charge, mz)
+        ccs = im2ccs(im, mz, charge)
 
         assert isinstance(ccs, np.ndarray)
         assert len(ccs) == len(im)
@@ -187,34 +185,34 @@ def test_ccs2im_im2ccs_roundtrip(self):
         charge = 2
         mz = 500.0
 
-        im = ccs2im(ccs_original, charge, mz)
-        ccs_roundtrip = im2ccs(im, charge, mz)
+        im = ccs2im(ccs_original, mz, charge)
+        ccs_roundtrip = im2ccs(im, mz, charge)
 
         assert abs(ccs_roundtrip - ccs_original) < 0.01
 
     def test_ccs2im_zero_values(self):
         """Test handling of zero values."""
         with pytest.raises((ValueError, ZeroDivisionError)):
-            ccs2im(0, 2, 500)
+            ccs2im(0, 500, 2)
 
     def test_im2ccs_zero_values(self):
         """Test handling of zero values."""
         with pytest.raises((ValueError, ZeroDivisionError)):
-            im2ccs(0, 2, 500)
+            im2ccs(0, 500, 2)
 
     def test_ccs2im_negative_values(self):
         """Test handling of negative values."""
         # Function should raise ValueError for negative CCS values
         with pytest.raises(ValueError, match="CCS must be positive"):
-            ccs2im(-450.0, 2, 500)
+            ccs2im(-450.0, 500, 2)
 
     def test_im2ccs_different_charges(self):
         """Test conversions with different charge states."""
         im = 1.0
         mz = 500
 
-        ccs_z2 = im2ccs(im, 2, mz)
-        ccs_z3 = im2ccs(im, 3, mz)
+        ccs_z2 = im2ccs(im, mz, 2)
+        ccs_z3 = im2ccs(im, mz, 3)
 
         assert ccs_z2 != ccs_z3
         assert ccs_z2 > 0 and ccs_z3 > 0

Original file line number	Diff line number	Diff line change
`@@ -130,6 +130,7 @@ def predict(`
`130`	`130`	`config=_get_model_config(multi=multi),`
`131`	`131`	`criterion=_get_loss_function(multi=multi),`
`132`	`132`	`)`
	`133`	`+ model.to(device)`
`133`	`134`	`model.eval()`
`134`	`135`	`LOGGER.debug(f"Model loaded on device: {device}")`
`135`	`136`