compomics / ms2rescore

Modular and user-friendly platform for AI-assisted rescoring of peptide identifications
https://ms2rescore.readthedocs.io
Apache License 2.0
39 stars 14 forks source link

Specify custom model directory in MS2PIPFeatureGenerator #134

Closed vrkosk closed 3 months ago

vrkosk commented 3 months ago

I'd like to specify a directory where I keep MS2PIP models, rather than the default ~/.ms2pip. This is currently impossible when using MS2PIPFeatureGenerator. ms2pip.correlate() accepts model_dir argument, but MS2PIPFeatureGenerator has no code to pass it through. So, ms2pip.correlate() gets model_dir=None and defaults to ~/.ms2pip internally.

Proposal is to add it as constructor argument:

class MS2PIPFeatureGenerator(FeatureGeneratorBase):
    """Generate MS²PIP-based features."""

    def __init__(
        self,
        *args,
        model: str = "HCD",
        model_dir: Optional[str] = None,  ## one-line change here
        ms2_tolerance: float = 0.02,

Set it as attribute:

        super().__init__(*args, **kwargs)
        self.model = model
        self.model_dir = model_dir  ## one-line change here

Then in add_features(), pass it to correlate():

                    ms2pip_results = correlate(
                        psms=psm_list_run,
                        spectrum_file=spectrum_filename,
                        spectrum_id_pattern=self.spectrum_id_pattern,
                        model=self.model,
                        model_dir=self.model_dir,  ## one-line change here