corems.molecular_id.search.database_interfaces

   1import os
   2import re
   3from abc import ABC
   4from io import StringIO
   5from pathlib import Path
   6
   7import numpy as np
   8import requests
   9import pandas as pd
  10from ms_entropy import FlashEntropySearch
  11
  12from corems.molecular_id.factory.EI_SQL import (
  13    EI_LowRes_SQLite,
  14    Metadatar,
  15    MetaboliteMetadata,
  16)
  17from corems.molecular_id.factory.lipid_molecular_metadata import LipidMetadata
  18from corems.mass_spectra.calc.lc_calc import find_closest
  19
  20
  21class SpectralDatabaseInterface(ABC):
  22    """
  23    Base class that facilitates connection to spectral reference databases,
  24    such as EMSL's Metabolomics Reference Database (MetabRef).
  25
  26    """
  27
  28    def __init__(self, key=None):
  29        """
  30        Initialize instance.
  31
  32        Parameters
  33        ----------
  34        key : str
  35            Token key.
  36
  37        """
  38
  39        self.key = key
  40
  41    def set_token(self, path):
  42        """
  43        Set environment variable for MetabRef database token.
  44
  45        Parameters
  46        ----------
  47        path : str
  48            Path to token.
  49
  50        """
  51
  52        # Read token from file
  53        with open(path, "r", encoding="utf-8") as f:
  54            token = f.readline().strip()
  55
  56        # Set environment variable
  57        os.environ[self.key] = token
  58
  59    def get_token(self):
  60        """
  61        Get environment variable for database token.
  62
  63        Returns
  64        -------
  65        str
  66            Token string.
  67
  68        """
  69
  70        # Check for token
  71        if self.key not in os.environ:
  72            raise ValueError("Must set {} environment variable.".format(self.key))
  73
  74        # Get token from environment variables
  75        return os.environ.get(self.key)
  76
  77    def get_header(self):
  78        """
  79        Access stored database token and prepare as header.
  80
  81        Returns
  82        -------
  83        str
  84            Header string.
  85
  86        """
  87
  88        # Get token
  89        token = self.get_token()
  90
  91        # Pad header information
  92        header = {"Authorization": f"Bearer {token}", "Content-Type": "text/plain"}
  93
  94        return header
  95
  96    def get_query(self, url, use_header=True):
  97        """
  98        Request payload from URL according to `get` protocol.
  99
 100        Parameters
 101        ----------
 102        url : str
 103            URL for request.
 104        use_header: bool
 105            Whether or not the query should include the header
 106
 107        Returns
 108        -------
 109        dict
 110            Response as JSON.
 111
 112        """
 113
 114        # Query URL via `get`
 115        if use_header:
 116            response = requests.get(url, headers=self.get_header())
 117        else:
 118            response = requests.get(url)
 119
 120        # Check response
 121        response.raise_for_status()
 122
 123        # Return as JSON
 124        return response.json()
 125
 126    def post_query(self, url, variable, values, tolerance):
 127        """
 128        Request payload from URL according to `post` protocol.
 129
 130        Parameters
 131        ----------
 132        url : str
 133            URL for request.
 134        variable : str
 135            Variable to query.
 136        values : str
 137            Specific values of `variable` to query.
 138        tolerance : str
 139            Query tolerance relative to `values`.
 140
 141        Returns
 142        -------
 143        dict
 144            Response as JSON.
 145
 146        """
 147
 148        # Coerce to string
 149        if not isinstance(variable, str):
 150            variable = str(variable).replace(" ", "")
 151
 152        if not isinstance(values, str):
 153            values = str(values).replace(" ", "")
 154
 155        if not isinstance(tolerance, str):
 156            tolerance = str(tolerance).replace(" ", "")
 157
 158        # Query URL via `post`
 159        response = requests.post(
 160            os.path.join(url, variable, tolerance),
 161            data=values,
 162            headers=self.get_header(),
 163        )
 164
 165        # Check response
 166        response.raise_for_status()
 167
 168        # Return as JSON
 169        return response.json()
 170
 171    def _check_flash_entropy_kwargs(self, fe_kwargs):
 172        """
 173        Check FlashEntropy keyword arguments.
 174
 175        Parameters
 176        ----------
 177        fe_kwargs : dict
 178            Keyword arguments for FlashEntropy search.
 179
 180
 181        Raises
 182        ------
 183        ValueError
 184            If "min_ms2_difference_in_da" or "max_ms2_tolerance_in_da" are present in `fe_kwargs` and they
 185            are not equal.
 186
 187        """
 188        # If "min_ms2_difference_in_da" in fe_kwargs, check that "max_ms2_tolerance_in_da" is also present and that min_ms2_difference_in_da = 2xmax_ms2_tolerance_in_da
 189        if (
 190            "min_ms2_difference_in_da" in fe_kwargs
 191            or "max_ms2_tolerance_in_da" in fe_kwargs
 192        ):
 193            if (
 194                "min_ms2_difference_in_da" not in fe_kwargs
 195                or "max_ms2_tolerance_in_da" not in fe_kwargs
 196            ):
 197                raise ValueError(
 198                    "Both 'min_ms2_difference_in_da' and 'max_ms2_tolerance_in_da' must be specified."
 199                )
 200            if (
 201                fe_kwargs["min_ms2_difference_in_da"]
 202                != 2 * fe_kwargs["max_ms2_tolerance_in_da"]
 203            ):
 204                raise ValueError(
 205                    "The values of 'min_ms2_difference_in_da' must be exactly 2x 'max_ms2_tolerance_in_da'."
 206                )
 207
 208    def _get_format_func(self, format):
 209        """
 210        Obtain format function by key.
 211
 212        Returns
 213        -------
 214        func
 215            Formatting function.
 216        """
 217
 218        if format.lower() in self.format_map.keys():
 219            return self.format_map[format.lower()]
 220
 221        raise ValueError(("{} not a supported format.").format(format))
 222
 223    def _dict_to_dataclass(self, metabref_lib, data_class):
 224        """
 225        Convert dictionary to dataclass.
 226
 227        Notes
 228        -----
 229        This function will pull the attributes a dataclass and its parent class
 230        and convert the dictionary to a dataclass instance with the appropriate
 231        attributes.
 232
 233        Parameters
 234        ----------
 235        data_class : :obj:`~dataclasses.dataclass`
 236            Dataclass to convert to.
 237        metabref_lib : dict
 238            Metabref dictionary object to convert to dataclass.
 239
 240        Returns
 241        -------
 242        :obj:`~dataclasses.dataclass`
 243            Dataclass instance.
 244
 245        """
 246
 247        # Get list of expected attributes of data_class
 248        data_class_keys = list(data_class.__annotations__.keys())
 249
 250        # Does the data_class inherit from another class, if so, get the attributes of the parent class as well
 251        if len(data_class.__mro__) > 2:
 252            parent_class_keys = list(data_class.__bases__[0].__annotations__.keys())
 253            data_class_keys = list(set(data_class_keys + parent_class_keys))
 254
 255        # Remove keys that are not in the data_class from the input dictionary
 256        input_dict = {k: v for k, v in metabref_lib.items() if k in data_class_keys}
 257
 258        # Add keys that are in the data class but not in the input dictionary as None
 259        for key in data_class_keys:
 260            if key not in input_dict.keys():
 261                input_dict[key] = None
 262        return data_class(**input_dict)
 263
 264    @staticmethod
 265    def normalize_peaks(arr):
 266        """
 267        Normalize peaks in an array.
 268
 269        Parameters
 270        ----------
 271        arr : :obj:`~numpy.array`
 272            Array of shape (N, 2), with m/z in the first column and abundance in
 273            the second.
 274
 275        Returns
 276        -------
 277        :obj:`~numpy.array`
 278            Normalized array of shape (N, 2), with m/z in the first column and
 279            normalized abundance in the second.
 280        """
 281        # Normalize the array
 282        arr[:, -1] = arr[:, -1] / arr[:, -1].sum()
 283
 284        return arr
 285
 286    @staticmethod
 287    def _build_flash_entropy_index(fe_lib, fe_kwargs={}, clean_spectra=True):
 288        """
 289        Build FlashEntropy index.
 290
 291        Parameters
 292        ----------
 293        fe_lib : list
 294            List of spectra to build index from. Can be a list of dictionaries or
 295            a FlashEntropy search instance.
 296        fe_kwargs : dict, optional
 297            Keyword arguments for FlashEntropy search.
 298        clean_spectra : bool, optional
 299            Clean spectra before building index. Default is True.
 300
 301        Returns
 302        -------
 303        :obj:`~ms_entropy.FlashEntropySearch`
 304            FlashEntropy search instance.
 305
 306        """
 307        # Initialize FlashEntropy
 308        fe_init_kws = [
 309            "max_ms2_tolerance_in_da",
 310            "mz_index_step",
 311            "low_memory",
 312            "path_data",
 313        ]
 314        fe_init_kws = {k: v for k, v in fe_kwargs.items() if k in fe_init_kws}
 315        fes = FlashEntropySearch(**fe_init_kws)
 316
 317        # Build FlashEntropy index
 318        fe_index_kws = [
 319            "max_indexed_mz",
 320            "precursor_ions_removal_da",
 321            "noise_threshold",
 322            "min_ms2_difference_in_da",
 323            "max_peak_num",
 324        ]
 325        fe_index_kws = {k: v for k, v in fe_kwargs.items() if k in fe_index_kws}
 326        fes.build_index(fe_lib, **fe_index_kws, clean_spectra=clean_spectra)
 327
 328        return fes
 329
 330
 331class MetabRefInterface(SpectralDatabaseInterface):
 332    """
 333    Interface to the Metabolomics Reference Database.
 334    """
 335
 336    def __init__(self):
 337        """
 338        Initialize instance.
 339
 340        """
 341
 342        super().__init__(key=None)
 343
 344    def spectrum_to_array(self, spectrum, normalize=True):
 345        """
 346        Convert MetabRef-formatted spectrum to array.
 347
 348        Parameters
 349        ----------
 350        spectrum : str
 351            MetabRef spectrum, i.e. list of (m/z,abundance) pairs.
 352        normalize : bool
 353            Normalize the spectrum by its magnitude.
 354
 355        Returns
 356        -------
 357        :obj:`~numpy.array`
 358            Array of shape (N, 2), with m/z in the first column and abundance in
 359            the second.
 360
 361        """
 362
 363        # Convert parenthesis-delimited string to array
 364        arr = np.array(
 365            re.findall(r"\(([^,]+),([^)]+)\)", spectrum), dtype=float
 366        ).reshape(-1, 2)
 367
 368        if normalize:
 369            arr = self.normalize_peaks(arr)
 370
 371        return arr
 372
 373    def _to_flashentropy(self, metabref_lib, normalize=True, fe_kwargs={}):
 374        """
 375        Convert metabref-formatted library to FlashEntropy library.
 376
 377        Parameters
 378        ----------
 379        metabref_lib : dict
 380            MetabRef MS2 library in JSON format or FlashEntropy search instance (for reformatting at different MS2 separation).
 381        normalize : bool
 382            Normalize each spectrum by its magnitude.
 383        fe_kwargs : dict, optional
 384            Keyword arguments for instantiation of FlashEntropy search and building index for FlashEntropy search;
 385            any keys not recognized will be ignored. By default, all parameters set to defaults.
 386
 387        Returns
 388        -------
 389        :obj:`~ms_entropy.FlashEntropySearch`
 390            MS2 library as FlashEntropy search instance.
 391
 392        Raises
 393        ------
 394        ValueError
 395            If "min_ms2_difference_in_da" or "max_ms2_tolerance_in_da" are present in `fe_kwargs` and they are not equal.
 396
 397        """
 398        self._check_flash_entropy_kwargs(fe_kwargs)
 399
 400        # Initialize empty library
 401        fe_lib = []
 402
 403        # Enumerate spectra
 404        for i, source in enumerate(metabref_lib):
 405            # Reorganize source dict, if necessary
 406            if "spectrum_data" in source.keys():
 407                spectrum = source["spectrum_data"]
 408            else:
 409                spectrum = source
 410
 411            # Rename precursor_mz key for FlashEntropy
 412            if "precursor_mz" not in spectrum.keys():
 413                spectrum["precursor_mz"] = spectrum.pop("precursor_ion")
 414
 415            # Convert CoreMS spectrum to array and clean, store as `peaks`
 416            spectrum["peaks"] = self.spectrum_to_array(
 417                spectrum["mz"], normalize=normalize
 418            )
 419
 420            # Add spectrum to library
 421            fe_lib.append(spectrum)
 422
 423        # Build FlashEntropy index
 424        fe_search = self._build_flash_entropy_index(fe_lib, fe_kwargs=fe_kwargs)
 425
 426        return fe_search
 427
 428    def get_query(self, url, use_header=False):
 429        """Overwrites the get_query method on the parent class to default to not use a header
 430
 431        Notes
 432        -----
 433        As of January 2025, the metabref database no longer requires a token and therefore no header is needed
 434
 435        """
 436        return super().get_query(url, use_header)
 437
 438
 439class MetabRefGCInterface(MetabRefInterface):
 440    """
 441    Interface to the Metabolomics Reference Database.
 442    """
 443
 444    def __init__(self):
 445        """
 446        Initialize instance.
 447
 448        """
 449
 450        super().__init__()
 451        self.GCMS_LIBRARY_URL = "https://metabref.emsl.pnnl.gov/api/mslevel/1"
 452        self.FAMES_URL = "https://metabref.emsl.pnnl.gov/api/fames"
 453
 454        self.__init_format_map__()
 455
 456    def __init_format_map__(self):
 457        """
 458        Initialize database format mapper, enabling multiple format requests.
 459
 460        """
 461
 462        # Define format workflows
 463        self.format_map = {
 464            "json": lambda x, normalize, fe_kwargs: x,
 465            "dict": lambda x,
 466            normalize,
 467            fe_kwargs: self._to_LowResolutionEICompound_dict(x, normalize),
 468            "sql": lambda x,
 469            normalize,
 470            fe_kwargs: self._LowResolutionEICompound_dict_to_sqlite(
 471                self._to_LowResolutionEICompound_dict(x, normalize)
 472            ),
 473        }
 474
 475        # Add aliases
 476        self.format_map["metabref"] = self.format_map["json"]
 477        self.format_map["datadict"] = self.format_map["dict"]
 478        self.format_map["data-dict"] = self.format_map["dict"]
 479        self.format_map["lowreseicompound"] = self.format_map["dict"]
 480        self.format_map["lowres"] = self.format_map["dict"]
 481        self.format_map["lowresgc"] = self.format_map["dict"]
 482        self.format_map["sqlite"] = self.format_map["sql"]
 483
 484    def available_formats(self):
 485        """
 486        View list of available formats.
 487
 488        Returns
 489        -------
 490        list
 491            Format map keys.
 492        """
 493
 494        return list(self.format_map.keys())
 495
 496    def get_library(self, format="json", normalize=False):
 497        """
 498        Request MetabRef GC/MS library.
 499
 500        Parameters
 501        ----------
 502        format : str
 503            Format of requested library, i.e. "json", "sql", "flashentropy".
 504            See `available_formats` method for aliases.
 505        normalize : bool
 506            Normalize the spectrum by its magnitude.
 507
 508        Returns
 509        -------
 510        Library in requested format.
 511
 512        """
 513
 514        # Init format function
 515        format_func = self._get_format_func(format)
 516
 517        return format_func(
 518            self.get_query(self.GCMS_LIBRARY_URL)["GC-MS"], normalize, {}
 519        )
 520
 521    def get_fames(self, format="json", normalize=False):
 522        """
 523        Request MetabRef GC/MS FAMEs library.
 524
 525        Parameters
 526        ----------
 527        format : str
 528            Format of requested library, i.e. "json", "sql", "flashentropy".
 529            See `available_formats` method for aliases.
 530        normalize : bool
 531            Normalize the spectrum by its magnitude.
 532
 533        Returns
 534        -------
 535        Library in requested format.
 536
 537        """
 538
 539        # Init format function
 540        format_func = self._get_format_func(format)
 541
 542        return format_func(self.get_query(self.FAMES_URL)["GC-MS"], normalize, {})
 543
 544    def _to_LowResolutionEICompound_dict(self, metabref_lib, normalize=False):
 545        """
 546        Convert MetabRef-formatted library to CoreMS LowResolutionEICompound-formatted
 547        dictionary for local ingestion.
 548
 549        Parameters
 550        ----------
 551        metabref_lib : dict
 552            MetabRef GC-MS library in JSON format.
 553        normalize : bool
 554            Normalize each spectrum by its magnitude.
 555
 556        Returns
 557        -------
 558        list of dict
 559            List of each spectrum contained in dictionary.
 560
 561        """
 562
 563        # All below key:value lookups are based on CoreMS class definitions
 564        # NOT MetabRef content. For example, MetabRef has keys for PubChem,
 565        # USI, etc. that are not considered below.
 566
 567        # Dictionary to map metabref keys to corems keys
 568        metadatar_cols = {
 569            "casno": "cas",
 570            "inchikey": "inchikey",
 571            "inchi": "inchi",
 572            "chebi": "chebi",
 573            "smiles": "smiles",
 574            "kegg": "kegg",
 575            "iupac_name": "iupac_name",
 576            "traditional_name": "traditional_name",  # Not present in metabref
 577            "common_name": "common_name",  # Not present in metabref
 578        }
 579
 580        # Dictionary to map metabref keys to corems keys
 581        lowres_ei_compound_cols = {
 582            "id": "metabref_id",
 583            "molecule_name": "name",  # Is this correct?
 584            "classify": "classify",  # Not present in metabref
 585            "formula": "formula",
 586            "ri": "ri",
 587            "rt": "retention_time",
 588            "source": "source",  # Not present in metabref
 589            "casno": "casno",
 590            "comments": "comment",
 591            "source_temp_c": "source_temp_c",  # Not present in metabref
 592            "ev": "ev",  # Not present in metabref
 593            "peak_count": "peaks_count",
 594            "mz": "mz",
 595            "abundance": "abundance",
 596        }
 597
 598        # Local result container
 599        corems_lib = []
 600
 601        # Enumerate spectra
 602        for i, source_ in enumerate(metabref_lib):
 603            # Copy source to prevent modification
 604            source = source_.copy()
 605
 606            # Flatten source dict
 607            source = source.pop("spectrum_data") | source
 608
 609            # Parse target data
 610            target = {
 611                lowres_ei_compound_cols[k]: v
 612                for k, v in source.items()
 613                if k in lowres_ei_compound_cols
 614            }
 615
 616            # Explicitly add this to connect with LowResCompoundRef later
 617            target["rt"] = source["rt"]
 618
 619            # Parse (mz, abundance)
 620            arr = self.spectrum_to_array(target["mz"], normalize=normalize)
 621            target["mz"] = arr[:, 0]
 622            target["abundance"] = arr[:, 1]
 623
 624            # Parse meta data
 625            target["metadata"] = {
 626                metadatar_cols[k]: v for k, v in source.items() if k in metadatar_cols
 627            }
 628
 629            # Add anything else
 630            for k in source:
 631                if k not in lowres_ei_compound_cols:
 632                    target[k] = source[k]
 633
 634            # Add to CoreMS list
 635            corems_lib.append(target)
 636
 637        return corems_lib
 638
 639    def _LowResolutionEICompound_dict_to_sqlite(
 640        self, lowres_ei_compound_dict, url="sqlite://"
 641    ):
 642        """
 643        Convert CoreMS LowResolutionEICompound-formatted dictionary to SQLite
 644        database for local ingestion.
 645
 646        Parameters
 647        ----------
 648        lowres_ei_compound_dict : dict
 649            CoreMS GC-MS library formatted for LowResolutionEICompound.
 650        url : str
 651            URL to SQLite prefix.
 652
 653        Returns
 654        -------
 655        sqlite database
 656            Spectra contained in SQLite database.
 657
 658        """
 659
 660        # Dictionary to map corems keys to all-caps keys
 661        capped_cols = {
 662            "name": "NAME",
 663            "formula": "FORM",
 664            "ri": "RI",
 665            "retention_time": "RT",
 666            "source": "SOURCE",
 667            "casno": "CASNO",
 668            "comment": "COMMENT",
 669            "peaks_count": "NUM PEAKS",
 670        }
 671
 672        # Initialize SQLite object
 673        sqlite_obj = EI_LowRes_SQLite(url=url)
 674
 675        # Iterate spectra
 676        for _data_dict in lowres_ei_compound_dict:
 677            # Copy source to prevent modification
 678            data_dict = _data_dict.copy()
 679
 680            # Add missing capped values
 681            for k, v in capped_cols.items():
 682                # Key exists
 683                if k in data_dict:
 684                    # # This will replace the key
 685                    # data_dict[v] = data_dict.pop(k)
 686
 687                    # This will keep both keys
 688                    data_dict[v] = data_dict[k]
 689
 690            # Parse number of peaks
 691            if not data_dict.get("NUM PEAKS"):
 692                data_dict["NUM PEAKS"] = len(data_dict.get("mz"))
 693
 694            # Parse CAS number
 695            if not data_dict.get("CASNO"):
 696                data_dict["CASNO"] = data_dict.get("CAS")
 697
 698            if not data_dict["CASNO"]:
 699                data_dict["CASNO"] = 0
 700
 701            # Build linked metadata table
 702            if "metadata" in data_dict:
 703                if len(data_dict["metadata"]) > 0:
 704                    data_dict["metadatar"] = Metadatar(**data_dict.pop("metadata"))
 705                else:
 706                    data_dict.pop("metadata")
 707
 708            # Attempt addition to sqlite
 709            try:
 710                sqlite_obj.add_compound(data_dict)
 711            except:
 712                print(data_dict["NAME"])
 713
 714        return sqlite_obj
 715
 716
 717class MetabRefLCInterface(MetabRefInterface):
 718    """
 719    Interface to the Metabolomics Reference Database for LC-MS data.
 720    """
 721
 722    def __init__(self):
 723        """
 724        Initialize instance.
 725
 726        """
 727
 728        super().__init__()
 729
 730        # API endpoint for precursor m/z search
 731        # inputs = mz, tolerance (in Da), polarity, page_no, per_page
 732        self.PRECURSOR_MZ_URL = "https://metabref.emsl.pnnl.gov/api/precursors/m/{}/t/{}/{}?page={}&per_page={}"
 733
 734        # API endpoint for returning full list of precursor m/z values in database
 735        # inputs = polarity, page_no, per_page
 736        self.PRECURSOR_MZ_ALL_URL = (
 737            "https://metabref.emsl.pnnl.gov/api/precursors/{}?page={}&per_page={}"
 738        )
 739
 740        self.__init_format_map__()
 741
 742    def __init_format_map__(self):
 743        """
 744        Initialize database format mapper, enabling multiple format requests.
 745
 746        """
 747
 748        # Define format workflows
 749        self.format_map = {
 750            "json": lambda x, normalize, fe_kwargs: x,
 751            "flashentropy": lambda x, normalize, fe_kwargs: self._to_flashentropy(
 752                x, normalize, fe_kwargs
 753            ),
 754        }
 755
 756        # Add aliases
 757        self.format_map["metabref"] = self.format_map["json"]
 758        self.format_map["fe"] = self.format_map["flashentropy"]
 759        self.format_map["flash-entropy"] = self.format_map["flashentropy"]
 760
 761    def query_by_precursor(
 762        self, mz_list, polarity, mz_tol_ppm, mz_tol_da_api=0.2, max_per_page=50
 763    ):
 764        """
 765        Query MetabRef by precursor m/z values.
 766
 767        Parameters
 768        ----------
 769        mz_list : list
 770            List of precursor m/z values.
 771        polarity : str
 772            Ionization polarity, either "positive" or "negative".
 773        mz_tol_ppm : float
 774            Tolerance in ppm for each precursor m/z value.
 775            Used for retrieving from a potential match from database.
 776        mz_tol_da_api : float, optional
 777            Maximum tolerance between precursor m/z values for API search, in daltons.
 778            Used to group similar mzs into a single API query for speed. Default is 0.2.
 779        max_per_page : int, optional
 780            Maximum records to return from MetabRef API query at a time.  Default is 50.
 781
 782        Returns
 783        -------
 784        list
 785            List of library entries in original JSON format.
 786        """
 787
 788        # If polarity is anything other than positive or negative, raise error
 789        if polarity not in ["positive", "negative"]:
 790            raise ValueError("Polarity must be 'positive' or 'negative'")
 791
 792        # Cluster groups of mz according to mz_tol_da_api for precursor query
 793        mz_list.sort()
 794        mz_groups = [[mz_list[0]]]
 795        for x in mz_list[1:]:
 796            if abs(x - mz_groups[-1][0]) <= mz_tol_da_api:
 797                mz_groups[-1].append(x)
 798            else:
 799                mz_groups.append([x])
 800
 801        # Query MetabRef for each mz group
 802        lib = []
 803        for mz_group in mz_groups:
 804            mz = np.mean(mz_group)
 805            if len(mz_group) == 1:
 806                mz = mz_group[0]
 807                tol = mz_tol_ppm * 10**-6 * mz
 808            else:
 809                mz = (max(mz_group) - min(mz_group)) / 2 + min(mz_group)
 810                tol = (max(mz_group) - min(mz_group)) / 2 + mz_tol_ppm**-6 * max(
 811                    mz_group
 812                )
 813
 814            # Get first page of results
 815            response = self.get_query(
 816                self.PRECURSOR_MZ_URL.format(
 817                    str(mz), str(tol), polarity, 1, max_per_page
 818                )
 819            )
 820            lib = lib + response["results"]
 821
 822            # If there are more pages of results, get them
 823            if response["total_pages"] > 1:
 824                for i in np.arange(2, response["total_pages"] + 1):
 825                    lib = (
 826                        lib
 827                        + self.get_query(
 828                            self.PRECURSOR_MZ_URL.format(
 829                                str(mz), str(tol), polarity, i, max_per_page
 830                            )
 831                        )["results"]
 832                    )
 833
 834        return lib
 835
 836    def request_all_precursors(self, polarity, per_page=50000):
 837        """
 838        Request all precursor m/z values for MS2 spectra from MetabRef.
 839
 840        Parameters
 841        ----------
 842        polarity : str
 843            Ionization polarity, either "positive" or "negative".
 844        per_page : int, optional
 845            Number of records to fetch per call. Default is 50000
 846
 847        Returns
 848        -------
 849        list
 850            List of all precursor m/z values, sorted.
 851        """
 852        # If polarity is anything other than positive or negative, raise error
 853        if polarity not in ["positive", "negative"]:
 854            raise ValueError("Polarity must be 'positive' or 'negative'")
 855
 856        precursors = []
 857
 858        # Get first page of results and total number of pages of results
 859        response = self.get_query(
 860            self.PRECURSOR_MZ_ALL_URL.format(polarity, str(1), str(per_page))
 861        )
 862        total_pages = response["total_pages"]
 863        precursors.extend([x["precursor_ion"] for x in response["results"]])
 864
 865        # Go through remaining pages of results
 866        for i in np.arange(2, total_pages + 1):
 867            response = self.get_query(
 868                self.PRECURSOR_MZ_ALL_URL.format(polarity, str(i), str(per_page))
 869            )
 870            precursors.extend([x["precursor_ion"] for x in response["results"]])
 871
 872        # Sort precursors from smallest to largest and remove duplicates
 873        precursors = list(set(precursors))
 874        precursors.sort()
 875
 876        return precursors
 877
 878    def get_lipid_library(
 879        self,
 880        mz_list,
 881        polarity,
 882        mz_tol_ppm,
 883        mz_tol_da_api=0.2,
 884        format="json",
 885        normalize=True,
 886        fe_kwargs={},
 887    ):
 888        """
 889        Request MetabRef lipid library.
 890
 891        Parameters
 892        ----------
 893        mz_list : list
 894            List of precursor m/z values.
 895        polarity : str
 896            Ionization polarity, either "positive" or "negative".
 897        mz_tol_ppm : float
 898            Tolerance in ppm for each precursor m/z value.
 899            Used for retrieving from a potential match from database.
 900        mz_tol_da_api : float, optional
 901            Maximum tolerance between precursor m/z values for API search, in daltons.
 902            Used to group similar mzs into a single API query for speed. Default is 0.2.
 903        format : str, optional
 904            Format of requested library, i.e. "json", "sql", "flashentropy".
 905            See `available_formats` method for aliases. Default is "json".
 906        normalize : bool, optional
 907            Normalize the spectrum by its magnitude. Default is True.
 908        fe_kwargs : dict, optional
 909            Keyword arguments for FlashEntropy search. Default is {}.
 910
 911        Returns
 912        -------
 913        tuple
 914            Library in requested format and lipid metadata as a LipidMetadata dataclass.
 915
 916        """
 917        mz_list.sort()
 918        mz_list = np.array(mz_list)
 919
 920        # Get all precursors in the library matching the polarity
 921        precusors_in_lib = self.request_all_precursors(polarity=polarity)
 922        precusors_in_lib = np.array(precusors_in_lib)
 923
 924        # Compare the mz_list with the precursors in the library, keep any mzs that are within mz_tol of any precursor in the library
 925        lib_mz_df = pd.DataFrame(precusors_in_lib, columns=["lib_mz"])
 926        lib_mz_df["closest_obs_mz"] = mz_list[
 927            find_closest(mz_list, lib_mz_df.lib_mz.values)
 928        ]
 929        lib_mz_df["mz_diff_ppm"] = np.abs(
 930            (lib_mz_df["lib_mz"] - lib_mz_df["closest_obs_mz"])
 931            / lib_mz_df["lib_mz"]
 932            * 1e6
 933        )
 934        lib_mz_sub = lib_mz_df[lib_mz_df["mz_diff_ppm"] <= mz_tol_ppm]
 935
 936        # Do the same in the opposite direction
 937        mz_df = pd.DataFrame(mz_list, columns=["mass_feature_mz"])
 938        mz_df["closest_lib_pre_mz"] = precusors_in_lib[
 939            find_closest(precusors_in_lib, mz_df.mass_feature_mz.values)
 940        ]
 941        mz_df["mz_diff_ppm"] = np.abs(
 942            (mz_df["mass_feature_mz"] - mz_df["closest_lib_pre_mz"])
 943            / mz_df["mass_feature_mz"]
 944            * 1e6
 945        )
 946        mz_df_sub = mz_df[mz_df["mz_diff_ppm"] <= mz_tol_ppm]
 947
 948        # Evaluate which is fewer mzs - lib_mz_sub or mz_df_sub and use that as the input for next step
 949        if len(lib_mz_sub) < len(mz_df_sub):
 950            mzs_to_query = lib_mz_sub.lib_mz.values
 951        else:
 952            mzs_to_query = mz_df_sub.mass_feature_mz.values
 953
 954        # Query the library for the precursors in the mz_list that are in the library to retrieve the spectra and metadata
 955        lib = self.query_by_precursor(
 956            mz_list=mzs_to_query,
 957            polarity=polarity,
 958            mz_tol_ppm=mz_tol_ppm,
 959            mz_tol_da_api=mz_tol_da_api,
 960        )
 961
 962        # Pull out lipid metadata from the metabref library and convert to LipidMetadata dataclass
 963        mol_data_dict = {x["id"]: x["Molecular Data"] for x in lib}
 964        lipid_lib = {x["id"]: x["Lipid Tree"] for x in lib if "Lipid Tree" in x.keys()}
 965        mol_data_dict = {k: {**v, **lipid_lib[k]} for k, v in mol_data_dict.items()}
 966        mol_data_dict = {
 967            k: self._dict_to_dataclass(v, LipidMetadata)
 968            for k, v in mol_data_dict.items()
 969        }
 970
 971        # Remove lipid metadata from the metabref library
 972        lib = [
 973            {k: v for k, v in x.items() if k not in ["Molecular Data", "Lipid Tree"]}
 974            for x in lib
 975        ]
 976        # Unpack the 'Lipid Fragments' key and the 'MSO Data" key from each entry
 977        for x in lib:
 978            if "Lipid Fragments" in x.keys():
 979                x.update(x.pop("Lipid Fragments"))
 980            if "MSO Data" in x.keys():
 981                x.update(x.pop("MSO Data"))
 982
 983        # Format the spectral library
 984        format_func = self._get_format_func(format)
 985        lib = format_func(lib, normalize=normalize, fe_kwargs=fe_kwargs)
 986        return (lib, mol_data_dict)
 987
 988
 989class MSPInterface(SpectralDatabaseInterface):
 990    """
 991    Interface to parse NIST MSP files
 992    """
 993
 994    def __init__(self, file_path):
 995        """
 996        Initialize instance.
 997
 998        Parameters
 999        ----------
1000        file_path : str
1001            Path to a local MSP file.
1002
1003        Attributes
1004        ----------
1005        file_path : str
1006            Path to the MSP file.
1007        _file_content : str
1008            Content of the MSP file.
1009        _data_frame : :obj:`~pandas.DataFrame`
1010            DataFrame of spectra from the MSP file with unaltered content.
1011        """
1012        super().__init__(key=None)
1013
1014        self.file_path = file_path
1015        if not os.path.exists(self.file_path):
1016            raise FileNotFoundError(
1017                f"File {self.file_path} does not exist. Please check the file path."
1018            )
1019        with open(self.file_path, "r") as f:
1020            self._file_content = f.read()
1021
1022        self._data_frame = self._read_msp_file()
1023        self.__init_format_map__()
1024
1025    def __init_format_map__(self):
1026        """
1027        Initialize database format mapper, enabling multiple format requests.
1028
1029        """
1030
1031        # x is a pandas dataframe similar to self._data_frame format
1032        # Define format workflows
1033        self.format_map = {
1034            "msp": lambda x, normalize, fe_kwargs: self._to_msp(x, normalize),
1035            "flashentropy": lambda x, normalize, fe_kwargs: self._to_flashentropy(
1036                x, normalize, fe_kwargs
1037            ),
1038            "df": lambda x, normalize, fe_kwargs: self._to_df(x, normalize),
1039        }
1040
1041        # Add aliases
1042        self.format_map["fe"] = self.format_map["flashentropy"]
1043        self.format_map["flash-entropy"] = self.format_map["flashentropy"]
1044        self.format_map["dataframe"] = self.format_map["df"]
1045        self.format_map["data-frame"] = self.format_map["df"]
1046
1047    def _read_msp_file(self):
1048        """
1049        Reads the MSP files into the pandas dataframe, and sort/remove zero intensity ions in MS/MS spectra.
1050
1051        Returns
1052        -------
1053        :obj:`~pandas.DataFrame`
1054            DataFrame of spectra from the MSP file, exacly as it is in the file (no sorting, filtering etc)
1055        """
1056        # If input_dataframe is provided, return it it
1057        spectra = []
1058        spectrum = {}
1059
1060        f = StringIO(self._file_content)
1061        for line in f:
1062            line = line.strip()
1063            if not line:
1064                continue  # Skip empty lines
1065
1066            # Handle metadata
1067            if ":" in line:
1068                key, value = line.split(":", 1)
1069                key = key.strip().lower()
1070                value = value.strip()
1071
1072                if key == "name":
1073                    # Save current spectrum and start a new one
1074                    if spectrum:
1075                        spectra.append(spectrum)
1076                    spectrum = {"name": value, "peaks": []}
1077                else:
1078                    spectrum[key] = value
1079
1080            # Handle peak data (assumed to start with a number)
1081            elif line[0].isdigit():
1082                peaks = line.split()
1083                m_z = float(peaks[0])
1084                intensity = float(peaks[1])
1085                spectrum["peaks"].append(([m_z, intensity]))
1086        # Save the last spectrum
1087        if spectrum:
1088            spectra.append(spectrum)
1089
1090        df = pd.DataFrame(spectra)
1091        for column in df.columns:
1092            if column != "peaks":  # Skip 'peaks' column
1093                try:
1094                    df[column] = pd.to_numeric(df[column], errors="raise")
1095                except:
1096                    pass
1097        return df
1098
1099    def _to_df(self, input_dataframe, normalize=True):
1100        """
1101        Convert MSP-derived library to FlashEntropy library. 
1102
1103        Parameters
1104        ----------
1105        input_dataframe : :obj:`~pandas.DataFrame`
1106            Input DataFrame containing MSP-formatted spectra.
1107        normalize : bool, optional
1108            Normalize each spectrum by its magnitude.
1109            Default is True.
1110
1111        Returns
1112        -------
1113        :obj:`~pandas.DataFrame`
1114            DataFrame of with desired normalization
1115        """
1116        if not normalize:
1117            return input_dataframe
1118        else:
1119            # Convert to dictionary
1120            db_dict = input_dataframe.to_dict(orient="records")
1121
1122            # Initialize empty library
1123            lib = []
1124
1125            # Enumerate spectra
1126            for i, source in enumerate(db_dict):
1127                spectrum = source
1128                # Check that spectrum["peaks"] exists
1129                if "peaks" not in spectrum.keys():
1130                    raise KeyError(
1131                        "MSP not interpretted correctly, 'peaks' key not found in spectrum, check _dataframe attribute."
1132                    )
1133
1134                # Convert spectrum["peaks"] to numpy array
1135                if not isinstance(spectrum["peaks"], np.ndarray):
1136                    spectrum["peaks"] = np.array(spectrum["peaks"])
1137
1138                # Normalize peaks, if requested
1139                if normalize:
1140                    spectrum["peaks"] = self.normalize_peaks(spectrum["peaks"])
1141                    spectrum["num peaks"] = len(spectrum["peaks"])
1142
1143                # Add spectrum to library
1144                lib.append(spectrum)
1145            
1146            # Convert to DataFrame
1147            df = pd.DataFrame(lib)
1148            return df
1149    
1150    def _to_flashentropy(self, input_dataframe, normalize=True, fe_kwargs={}):
1151        """
1152        Convert MSP-derived library to FlashEntropy library.
1153
1154        Parameters
1155        ----------
1156        input_dataframe : :obj:`~pandas.DataFrame`
1157            Input DataFrame containing MSP-formatted spectra.
1158        normalize : bool
1159            Normalize each spectrum by its magnitude.
1160        fe_kwargs : dict, optional
1161            Keyword arguments for instantiation of FlashEntropy search and building index for FlashEntropy search;
1162            any keys not recognized will be ignored. By default, all parameters set to defaults.
1163
1164        Returns
1165        -------
1166        :obj:`~ms_entropy.FlashEntropySearch`
1167            MS2 library as FlashEntropy search instance.
1168
1169        Raises
1170        ------
1171        ValueError
1172            If "min_ms2_difference_in_da" or "max_ms2_tolerance_in_da" are present in `fe_kwargs` and they
1173        """
1174        self._check_flash_entropy_kwargs(fe_kwargs)
1175
1176        db_df = input_dataframe
1177
1178        # Convert to dictionary
1179        db_dict = db_df.to_dict(orient="records")
1180
1181        # Initialize empty library
1182        fe_lib = []
1183
1184        # Enumerate spectra
1185        for i, source in enumerate(db_dict):
1186            # Reorganize source dict, if necessary
1187            if "spectrum_data" in source.keys():
1188                spectrum = source["spectrum_data"]
1189            else:
1190                spectrum = source
1191
1192            # Rename precursor_mz key for FlashEntropy
1193            if "precursor_mz" not in spectrum.keys():
1194                if "precursormz" in spectrum:
1195                    spectrum["precursor_mz"] = spectrum.pop("precursormz")
1196                elif "precursor_ion" in spectrum:
1197                    spectrum["precursor_mz"] = spectrum.pop("precursor_ion")
1198                else:
1199                    raise KeyError(
1200                        "MSP must have either 'precursormz' or 'precursor_ion' key to be converted to FlashEntropy format."
1201                    )
1202
1203            # Check that spectrum["peaks"] exists
1204            if "peaks" not in spectrum.keys():
1205                raise KeyError(
1206                    "MSP not interpretted correctly, 'peaks' key not found in spectrum, check _dataframe attribute."
1207                )
1208
1209            # Convert spectrum["peaks"] to numpy array
1210            if not isinstance(spectrum["peaks"], np.ndarray):
1211                spectrum["peaks"] = np.array(spectrum["peaks"])
1212
1213            # Normalize peaks, if requested
1214            if normalize:
1215                spectrum["peaks"] = self.normalize_peaks(spectrum["peaks"])
1216
1217            # Add spectrum to library
1218            fe_lib.append(spectrum)
1219
1220        # Build FlashEntropy index
1221        fe_search = self._build_flash_entropy_index(fe_lib, fe_kwargs=fe_kwargs)
1222
1223        return fe_search
1224    
1225    def _check_msp_compatibility(self):
1226        """
1227        Check if the MSP file is compatible with the get_metabolomics_spectra_library method and provide feedback if it is not.
1228        """
1229        # Check polarity
1230        if (
1231            "polarity" not in self._data_frame.columns
1232            and "ionmode" not in self._data_frame.columns
1233        ):
1234            raise ValueError(
1235                "Neither 'polarity' nor 'ionmode' columns found in the input MSP metadata. Please check the file."
1236            )
1237        polarity_column = (
1238            "polarity" if "polarity" in self._data_frame.columns else "ionmode"
1239        )
1240
1241        # Check if polarity_column contents is either "positive" or "negative"
1242        if not all(self._data_frame[polarity_column].isin(["positive", "negative"])):
1243            raise ValueError(
1244                f"Input field on MSP '{polarity_column}' must contain only 'positive' or 'negative' values."
1245            )
1246
1247        # Check if the MSP file contains the required columns for metabolite metadata
1248        # inchikey, by name, not null
1249        # either formula or molecular_formula, not null
1250        if not all(self._data_frame["inchikey"].notnull()):
1251            raise ValueError(
1252                "Input field on MSP 'inchikey' must contain only non-null values."
1253            )
1254        if (
1255            "formula" not in self._data_frame.columns
1256            and "molecular_formula" not in self._data_frame.columns
1257        ):
1258            raise ValueError(
1259                "Input field on MSP must contain either 'formula' or 'molecular_formula' columns."
1260            )
1261        molecular_formula_column = (
1262            "formula" if "formula" in self._data_frame.columns else "molecular_formula"
1263        )
1264        if not all(self._data_frame[molecular_formula_column].notnull()):
1265            raise ValueError(
1266                f"Input field on MSP '{molecular_formula_column}' must contain only non-null values."
1267            )
1268
1269    def get_metabolomics_spectra_library(
1270        self,
1271        polarity,
1272        metabolite_metadata_mapping={},
1273        format="fe",
1274        normalize=True,
1275        fe_kwargs={},
1276    ):
1277        """
1278        Prepare metabolomics spectra library and associated metabolite metadata
1279
1280        Note: this uses the inchikey as the index for the metabolite metadata dataframe and for connecting to the spectra, so it must be in the input
1281
1282        """
1283        # Check if the MSP file is compatible with the get_metabolomics_spectra_library method
1284        self._check_msp_compatibility()
1285
1286        # Check if the polarity parameter is valid and if a polarity column exists in the dataframe
1287        if polarity not in ["positive", "negative"]:
1288            raise ValueError("Polarity must be 'positive' or 'negative'")
1289        polarity_column = (
1290            "polarity" if "polarity" in self._data_frame.columns else "ionmode"
1291        )
1292
1293        # Get a subset of the initial dataframea by polarity
1294        db_df = self._data_frame[self._data_frame[polarity_column] == polarity].copy()
1295
1296        # Rename the columns of the db_df to match the MetaboliteMetadata dataclass using the metabolite_metadata_mapping
1297        # If the mapping is not provided, use the default mapping
1298        if not metabolite_metadata_mapping:
1299            metabolite_metadata_mapping = {
1300                "chebi_id": "chebi",
1301                "kegg_id": "kegg",
1302                "refmet_name": "common_name",
1303                "molecular_formula": "formula",
1304                "gnps_spectra_id":"id",
1305                "precursormz": "precursor_mz",
1306                "precursortype":"ion_type"
1307            }
1308        db_df.rename(columns=metabolite_metadata_mapping, inplace=True)
1309        db_df["molecular_data_id"] = db_df["inchikey"]
1310
1311
1312
1313        # Check if the resulting dataframe has the required columns for the flash entropy search
1314        required_columns = ["molecular_data_id", "precursor_mz", "ion_type", "id"]
1315        for col in required_columns:
1316            if col not in db_df.columns:
1317                raise ValueError(
1318                    f"Input field on MSP must contain '{col}' column for FlashEntropy search."
1319                )
1320
1321        # Pull out the metabolite metadata from the dataframe and put it into a different dataframe
1322        # First get a list of the possible attributes of the MetaboliteMetadata dataclass
1323        metabolite_metadata_keys = list(MetaboliteMetadata.__annotations__.keys())
1324        # Replace id with molecular_data_id in metabolite_metadata_keys
1325        metabolite_metadata_keys = [
1326            "molecular_data_id" if x == "id" else x for x in metabolite_metadata_keys
1327        ]
1328        metabolite_metadata_df = db_df[
1329            db_df.columns[db_df.columns.isin(metabolite_metadata_keys)]
1330        ].copy()
1331
1332        # Make unique and recast the id column for metabolite metadata
1333        metabolite_metadata_df.drop_duplicates(subset=["molecular_data_id"], inplace=True)
1334        metabolite_metadata_df["id"] = metabolite_metadata_df["molecular_data_id"]
1335
1336        # Convert to a dictionary using the inchikey as the key
1337        metabolite_metadata_dict = metabolite_metadata_df.to_dict(
1338            orient="records"
1339        )
1340        metabolite_metadata_dict = {
1341            v["id"]: self._dict_to_dataclass(v, MetaboliteMetadata)
1342            for v in metabolite_metadata_dict
1343        }
1344
1345        # Remove the metabolite metadata columns from the original dataframe
1346        for key in metabolite_metadata_keys:
1347            if key != "molecular_data_id":
1348                if key in db_df.columns:
1349                    db_df.drop(columns=key, inplace=True)
1350
1351        # Format the spectral library
1352        format_func = self._get_format_func(format)
1353        lib = format_func(db_df, normalize=normalize, fe_kwargs=fe_kwargs)
1354        return (lib, metabolite_metadata_dict)