Merge pull request #237 from MannLabs/update-adaptive-frag

Update adaptive frag types interfaces

Author: mo-sameh

peptdeep/model/ms2.py

@@ -656,9 +656,30 @@ def predict(
         batch_size=1024,
         verbose=False,
         reference_frag_df=None,
+        allow_unsafe_predictions=False,
         **kwargs,
     ) -> pd.DataFrame:
-        if not self._safe_to_predict:
+        """
+        Predict MS2 fragment intensities
+
+        Parameters
+        ----------
+        precursor_df : pd.DataFrame
+            Precursor DataFrame
+        batch_size : int, optional
+            Batch size, by default 1024
+        verbose : bool, optional
+            Verbose, by default False
+        reference_frag_df : pd.DataFrame, optional
+            Reference fragment intensity DataFrame, by default None
+        allow_unsafe_predictions : bool, optional
+            Allow a newly a randomly initialized model to be used for prediction, by default False
+        Returns
+        -------
+        pd.DataFrame
+            Predicted fragment intensities
+        """
+        if not self._safe_to_predict and not allow_unsafe_predictions:
             raise ValueError(
                 f"The model is not safe to use for prediction. This might mean that the requested charged_frag_types {self.charged_frag_types} are not a subset of the charged_frag_types used to train the loaded pretrained model {self.model.supported_charged_frag_types}. Please choose a subset of the supported charged_frag_types or retrain the model with the requested charged_frag_types."
             )

peptdeep/pretrained_models.py

@@ -321,6 +321,20 @@ def __init__(
 
         self.reset_by_global_settings(reload_models=False)
 
+    def reinitialize_ms2_model(self, charged_frag_types: typing.List[str], **kwargs):
+        """
+        Reinitialize the MS2 model with new charged fragment types.
+
+        Parameters
+        ----------
+        charged_frag_types : List[str]
+            Charged fragment types for the new MS2 model.
+
+        kwargs : dict
+            Other keyword arguments for `pDeepModel`.
+        """
+        self.ms2_model = pDeepModel(charged_frag_types=charged_frag_types, **kwargs)
+
     def reset_by_global_settings(
         self,
         reload_models=True,
@@ -965,7 +979,7 @@ def mp_param_generator(df_groupby):
             fragment_mz_df_list = None
 
         if self.verbose:
-            logging.info(f'Predicting {",".join(predict_items)} ...')
+            logging.info(f"Predicting {','.join(predict_items)} ...")
         verbose_bak = self.verbose
         self.verbose = False