corems.molecular_id.search.molecularFormulaSearch

   1__author__ = "Yuri E. Corilo"
   2__date__ = "Jul 29, 2019"
   3
   4
   5from typing import List
   6
   7import tqdm
   8
   9from corems import chunks, timeit
  10from corems.encapsulation.constant import Labels
  11from corems.molecular_formula.factory.MolecularFormulaFactory import (
  12    LCMSLibRefMolecularFormula,
  13    MolecularFormula,
  14)
  15from corems.molecular_id.factory.MolecularLookupTable import MolecularCombinations
  16from corems.molecular_id.factory.molecularSQL import MolForm_SQL
  17from corems.ms_peak.factory.MSPeakClasses import _MSPeak
  18
  19last_error = 0
  20last_dif = 0
  21closest_error = 0
  22error_average = 0
  23nbValues = 0
  24
  25
  26class SearchMolecularFormulas:
  27    """Class for searching molecular formulas in a mass spectrum.
  28
  29    Parameters
  30    ----------
  31    mass_spectrum_obj : MassSpectrum
  32        The mass spectrum object.
  33    sql_db : MolForm_SQL, optional
  34        The SQL database object, by default None.
  35    first_hit : bool, optional
  36        Flag to indicate whether to skip peaks that already have a molecular formula assigned, by default False.
  37    find_isotopologues : bool, optional
  38        Flag to indicate whether to find isotopologues, by default True.
  39
  40    Attributes
  41    ----------
  42    mass_spectrum_obj : MassSpectrum
  43        The mass spectrum object.
  44    sql_db : MolForm_SQL
  45        The SQL database object.
  46    first_hit : bool
  47        Flag to indicate whether to skip peaks that already have a molecular formula assigned.
  48    find_isotopologues : bool
  49        Flag to indicate whether to find isotopologues.
  50
  51
  52    Methods
  53    -------
  54    * run_search().
  55        Run the molecular formula search.
  56    * run_worker_mass_spectrum().
  57        Run the molecular formula search on the mass spectrum object.
  58    * run_worker_ms_peaks().
  59        Run the molecular formula search on the given list of mass spectrum peaks.
  60    * database_to_dict().
  61        Convert the database results to a dictionary.
  62    * run_molecular_formula().
  63        Run the molecular formula search on the given list of mass spectrum peaks.
  64    * search_mol_formulas().
  65        Search for molecular formulas in the mass spectrum.
  66
  67    """
  68
  69    def __init__(
  70        self,
  71        mass_spectrum_obj,
  72        sql_db=None,
  73        first_hit: bool = False,
  74        find_isotopologues: bool = True,
  75    ):
  76        self.first_hit = first_hit
  77
  78        self.find_isotopologues = find_isotopologues
  79
  80        self.mass_spectrum_obj = mass_spectrum_obj
  81
  82        if not sql_db:
  83            self.sql_db = MolForm_SQL(
  84                url=mass_spectrum_obj.molecular_search_settings.url_database
  85            )
  86
  87        else:
  88            self.sql_db = sql_db
  89
  90    def __enter__(self):
  91        """Open the SQL database connection."""
  92        return self
  93
  94    def __exit__(self, exc_type, exc_val, exc_tb):
  95        """Close the SQL database connection."""
  96        self.sql_db.close()
  97
  98        return False
  99
 100    def run_search(
 101        self,
 102        mspeaks: list,
 103        query: dict,
 104        min_abundance: float,
 105        ion_type: str,
 106        ion_charge: int,
 107        adduct_atom=None,
 108    ):
 109        """Run the molecular formula search.
 110
 111        Parameters
 112        ----------
 113        mspeaks : list of MSPeak
 114            The list of mass spectrum peaks.
 115        query : dict
 116            The query dictionary containing the possible molecular formulas.
 117        min_abundance : float
 118            The minimum abundance threshold.
 119        ion_type : str
 120            The ion type.
 121        ion_charge : int
 122            The ion charge.
 123        adduct_atom : str, optional
 124            The adduct atom, by default None.
 125        """
 126
 127        def get_formulas(nominal_overlay: float = 0.1):
 128            """
 129            Get the list of formulas based on the nominal overlay.
 130
 131            Parameters
 132            ----------
 133            nominal_overlay : float, optional
 134                The nominal overlay, by default 0.1.
 135
 136            Returns
 137            -------
 138            list
 139                The list of formulas.
 140            """
 141            nominal_mz = ms_peak.nominal_mz_exp
 142
 143            defect_mass = ms_peak.mz_exp - nominal_mz
 144            nominal_masses = [nominal_mz]
 145
 146            if (defect_mass) >= 1 - nominal_overlay:
 147                nominal_masses.append(nominal_mz + 1)
 148            elif (defect_mass) <= nominal_overlay:
 149                nominal_masses.append(nominal_mz - 1)
 150
 151            list_formulas_candidates = []
 152
 153            for nominal_mass in nominal_masses:
 154                if nominal_mass in query.keys():
 155                    list_formulas_candidates.extend(query.get(nominal_mass))
 156
 157            return list_formulas_candidates
 158
 159        all_assigned_indexes = list()
 160
 161        # molecular_search_settings = self.mass_spectrum_obj.molecular_search_settings
 162
 163        search_molfrom = SearchMolecularFormulaWorker(
 164            find_isotopologues=self.find_isotopologues
 165        )
 166
 167        for ms_peak in mspeaks:
 168            # already assigned a molecular formula
 169            if self.first_hit:
 170                if ms_peak.is_assigned:
 171                    continue
 172
 173            ms_peak_indexes = search_molfrom.find_formulas(
 174                get_formulas(),
 175                min_abundance,
 176                self.mass_spectrum_obj,
 177                ms_peak,
 178                ion_type,
 179                ion_charge,
 180                adduct_atom,
 181            )
 182
 183            all_assigned_indexes.extend(ms_peak_indexes)
 184
 185        # all_assigned_indexes = MolecularFormulaSearchFilters().filter_isotopologue(all_assigned_indexes, self.mass_spectrum_obj)
 186
 187        # all_assigned_indexes = MolecularFormulaSearchFilters().filter_kendrick(all_assigned_indexes, self.mass_spectrum_obj)
 188
 189        # MolecularFormulaSearchFilters().check_min_peaks(all_assigned_indexes, self.mass_spectrum_obj)
 190        # filter per min peaks per mono isotopic class
 191
 192    def run_worker_mass_spectrum(self):
 193        """Run the molecular formula search on the mass spectrum object."""
 194        self.run_molecular_formula(
 195            self.mass_spectrum_obj.sort_by_abundance(), 
 196            print_time=self.mass_spectrum_obj.molecular_search_settings.verbose_processing
 197            )
 198
 199    def run_worker_ms_peaks(self, ms_peaks):
 200        """Run the molecular formula search on the given list of mass spectrum peaks.
 201
 202        Parameters
 203        ----------
 204        ms_peaks : list of MSPeak
 205            The list of mass spectrum peaks.
 206        """
 207        self.run_molecular_formula(
 208            ms_peaks,
 209            print_time=self.mass_spectrum_obj.molecular_search_settings.verbose_processing
 210        )
 211
 212    @staticmethod
 213    def database_to_dict(
 214        classe_str_list,
 215        nominal_mzs,
 216        mf_search_settings,
 217        ion_charge,
 218        sql_db=None,
 219    ):
 220        """Convert the database results to a dictionary.
 221
 222        Parameters
 223        ----------
 224        classe_str_list : list
 225            The list of class strings.
 226        nominal_mzs : list
 227            The list of nominal m/z values.
 228        mf_search_settings : MolecularFormulaSearchSettings
 229            The molecular formula search settings.
 230        ion_charge : int
 231            The ion charge.
 232        sql_db : MolForm_SQL, optional
 233            The SQL database object, by default None. If None, a new MolForm_SQL object will be created.
 234
 235        Returns
 236        -------
 237        dict
 238            The dictionary containing the database results.
 239        """
 240        owns_db = sql_db is None
 241        if owns_db:
 242            sql_db = MolForm_SQL(url=mf_search_settings.url_database)
 243
 244        try:
 245            dict_res = {}
 246
 247            if mf_search_settings.isProtonated:
 248                dict_res[Labels.protonated_de_ion] = sql_db.get_dict_by_classes(
 249                    classe_str_list,
 250                    Labels.protonated_de_ion,
 251                    nominal_mzs,
 252                    ion_charge,
 253                    mf_search_settings,
 254                )
 255
 256            if mf_search_settings.isRadical:
 257                dict_res[Labels.radical_ion] = sql_db.get_dict_by_classes(
 258                    classe_str_list,
 259                    Labels.radical_ion,
 260                    nominal_mzs,
 261                    ion_charge,
 262                    mf_search_settings,
 263                )
 264
 265            if mf_search_settings.isAdduct:
 266                adduct_list = (
 267                    mf_search_settings.adduct_atoms_neg
 268                    if ion_charge < 0
 269                    else mf_search_settings.adduct_atoms_pos
 270                )
 271                dict_res[Labels.adduct_ion] = sql_db.get_dict_by_classes(
 272                    classe_str_list,
 273                    Labels.adduct_ion,
 274                    nominal_mzs,
 275                    ion_charge,
 276                    mf_search_settings,
 277                    adducts=adduct_list,
 278                )
 279
 280            return dict_res
 281        finally:
 282            if owns_db:
 283                sql_db.close()
 284
 285    @timeit(print_time=True)
 286    def run_molecular_formula(self, ms_peaks, **kwargs):
 287        """Run the molecular formula search on the given list of mass spectrum peaks.
 288
 289        Parameters
 290        ----------
 291        ms_peaks : list of MSPeak
 292            The list of mass spectrum peaks.
 293        **kwargs
 294            Additional keyword arguments. 
 295            Most notably, print_time, which is a boolean flag to indicate whether to print the time 
 296            and passed to the timeit decorator.
 297        """
 298        ion_charge = self.mass_spectrum_obj.polarity
 299        min_abundance = self.mass_spectrum_obj.min_abundance
 300        nominal_mzs = self.mass_spectrum_obj.nominal_mz
 301
 302        verbose = self.mass_spectrum_obj.molecular_search_settings.verbose_processing
 303        # reset average error, only relevant is average mass error method is being used
 304        SearchMolecularFormulaWorker(
 305            find_isotopologues=self.find_isotopologues
 306        ).reset_error(self.mass_spectrum_obj)
 307
 308        # check database for all possible molecular formula combinations based on the setting passed to self.mass_spectrum_obj.molecular_search_settings
 309        classes = MolecularCombinations(self.sql_db).runworker(
 310            self.mass_spectrum_obj.molecular_search_settings,
 311            print_time=self.mass_spectrum_obj.molecular_search_settings.verbose_processing
 312        )
 313
 314        # split the database load to not blowout the memory
 315        # TODO add to the settings
 316        for classe_chunk in chunks(
 317            classes, self.mass_spectrum_obj.molecular_search_settings.db_chunk_size
 318        ):
 319            classes_str_list = [class_tuple[0] for class_tuple in classe_chunk]
 320
 321            # load the molecular formula objs binned by ion type and heteroatoms classes, {ion type:{classe:[list_formula]}}
 322            # for adduct ion type a third key is added {atoms:{ion type:{classe:[list_formula]}}}
 323            dict_res = self.database_to_dict(
 324                classes_str_list,
 325                nominal_mzs,
 326                self.mass_spectrum_obj.molecular_search_settings,
 327                ion_charge,
 328                sql_db=self.sql_db,
 329            )
 330            pbar = tqdm.tqdm(classe_chunk, disable = not verbose)
 331            for classe_tuple in pbar:
 332                # class string is a json serialized dict
 333                classe_str = classe_tuple[0]
 334                classe_dict = classe_tuple[1]
 335
 336                if self.mass_spectrum_obj.molecular_search_settings.isProtonated:
 337                    ion_type = Labels.protonated_de_ion
 338                    if verbose:
 339                        pbar.set_description_str(
 340                            desc="Started molecular formula search for class %s, (de)protonated "
 341                            % classe_str,
 342                            refresh=True,
 343                        )
 344
 345                    candidate_formulas = dict_res.get(ion_type).get(classe_str)
 346
 347                    if candidate_formulas:
 348                        self.run_search(
 349                            ms_peaks,
 350                            candidate_formulas,
 351                            min_abundance,
 352                            ion_type,
 353                            ion_charge,
 354                        )
 355
 356                if self.mass_spectrum_obj.molecular_search_settings.isRadical:
 357                    if verbose:
 358                        pbar.set_description_str(
 359                            desc="Started molecular formula search for class %s, radical "
 360                            % classe_str,
 361                            refresh=True,
 362                        )
 363
 364                    ion_type = Labels.radical_ion
 365
 366                    candidate_formulas = dict_res.get(ion_type).get(classe_str)
 367
 368                    if candidate_formulas:
 369                        self.run_search(
 370                            ms_peaks,
 371                            candidate_formulas,
 372                            min_abundance,
 373                            ion_type,
 374                            ion_charge,
 375                        )
 376                # looks for adduct, used_atom_valences should be 0
 377                # this code does not support H exchance by halogen atoms
 378                if self.mass_spectrum_obj.molecular_search_settings.isAdduct:
 379                    if verbose:
 380                        pbar.set_description_str(
 381                            desc="Started molecular formula search for class %s, adduct "
 382                            % classe_str,
 383                            refresh=True,
 384                        )
 385
 386                    ion_type = Labels.adduct_ion
 387                    dict_atoms_formulas = dict_res.get(ion_type)
 388
 389                    for adduct_atom, dict_by_class in dict_atoms_formulas.items():
 390                        candidate_formulas = dict_by_class.get(classe_str)
 391
 392                        if candidate_formulas:
 393                            self.run_search(
 394                                ms_peaks,
 395                                candidate_formulas,
 396                                min_abundance,
 397                                ion_type,
 398                                ion_charge,
 399                                adduct_atom=adduct_atom,
 400                            )
 401        self.sql_db.close()
 402
 403    def search_mol_formulas(
 404        self,
 405        possible_formulas_list: List[MolecularFormula],
 406        ion_type: str,
 407        neutral_molform=True,
 408        find_isotopologues=True,
 409        adduct_atom=None,
 410    ) -> List[_MSPeak]:
 411        """Search for molecular formulas in the mass spectrum.
 412
 413        Parameters
 414        ----------
 415        possible_formulas_list : list of MolecularFormula
 416            The list of possible molecular formulas.
 417        ion_type : str
 418            The ion type.
 419        neutral_molform : bool, optional
 420            Flag to indicate whether the molecular formulas are neutral, by default True.
 421        find_isotopologues : bool, optional
 422            Flag to indicate whether to find isotopologues, by default True.
 423        adduct_atom : str, optional
 424            The adduct atom, by default None.
 425
 426        Returns
 427        -------
 428        list of MSPeak
 429            The list of mass spectrum peaks with assigned molecular formulas.
 430        """
 431        # neutral_molform: some reference files already present the formula on ion mode, for instance, bruker reference files
 432        #    if that is the case than turn neutral_molform off
 433
 434        SearchMolecularFormulaWorker(find_isotopologues=find_isotopologues).reset_error(
 435            self.mass_spectrum_obj
 436        )
 437
 438        initial_min_peak_bool = (
 439            self.mass_spectrum_obj.molecular_search_settings.use_min_peaks_filter
 440        )
 441        initial_runtime_kendrick_filter = (
 442            self.mass_spectrum_obj.molecular_search_settings.use_runtime_kendrick_filter
 443        )
 444
 445        # Are the following 3 lines redundant?
 446        self.mass_spectrum_obj.molecular_search_settings.use_min_peaks_filter = False
 447        self.mass_spectrum_obj.molecular_search_settings.use_min_peaks_filter = (
 448            False  # TODO check this line
 449        )
 450        self.mass_spectrum_obj.molecular_search_settings.use_runtime_kendrick_filter = (
 451            False
 452        )
 453
 454        possible_formulas_dict_nm = {}
 455
 456        for mf in possible_formulas_list:
 457            if neutral_molform:
 458                nm = int(mf.protonated_mz)
 459            else:
 460                nm = int(mf.mz_nominal_calc)
 461
 462            if nm in possible_formulas_dict_nm.keys():
 463                possible_formulas_dict_nm[nm].append(mf)
 464
 465            else:
 466                possible_formulas_dict_nm[nm] = [mf]
 467
 468        min_abundance = self.mass_spectrum_obj.min_abundance
 469
 470        ion_type = ion_type
 471
 472        self.run_search(
 473            self.mass_spectrum_obj,
 474            possible_formulas_dict_nm,
 475            min_abundance,
 476            ion_type,
 477            self.mass_spectrum_obj.polarity,
 478            adduct_atom=adduct_atom,
 479        )
 480
 481        self.mass_spectrum_obj.molecular_search_settings.use_min_peaks_filter = (
 482            initial_min_peak_bool
 483        )
 484        self.mass_spectrum_obj.molecular_search_settings.use_runtime_kendrick_filter = (
 485            initial_runtime_kendrick_filter
 486        )
 487
 488        mspeaks = [mspeak for mspeak in self.mass_spectrum_obj if mspeak.is_assigned]
 489
 490        self.sql_db.close()
 491
 492        return mspeaks
 493
 494
 495class SearchMolecularFormulaWorker:
 496    """Class for searching molecular formulas in a mass spectrum.
 497
 498    Parameters
 499    ----------
 500    find_isotopologues : bool, optional
 501        Flag to indicate whether to find isotopologues, by default True.
 502
 503    Attributes
 504    ----------
 505    find_isotopologues : bool
 506        Flag to indicate whether to find isotopologues.
 507
 508    Methods
 509    -------
 510    * reset_error().
 511        Reset the error variables.
 512    * set_last_error().
 513        Set the last error.
 514    * find_formulas().
 515        Find the formulas.
 516    * calc_error().
 517        Calculate the error.
 518    """
 519
 520    # TODO add reset error function
 521    # needs this wraper to pass the class to multiprocessing
 522
 523    def __init__(self, find_isotopologues=True):
 524        self.find_isotopologues = find_isotopologues
 525
 526    def __call__(self, args):
 527        """Call the find formulas function.
 528
 529        Parameters
 530        ----------
 531        args : tuple
 532            The arguments.
 533
 534        Returns
 535        -------
 536        list
 537            The list of mass spectrum peaks with assigned molecular formulas.
 538        """
 539        return self.find_formulas(*args)  # ,args[1]
 540
 541    def reset_error(self, mass_spectrum_obj):
 542        """Reset the error variables.
 543
 544        Parameters
 545        ----------
 546        mass_spectrum_obj : MassSpectrum
 547            The mass spectrum object.
 548
 549        Notes
 550        -----
 551        This function resets the error variables for the given mass spectrum object.
 552        """
 553        global last_error, last_dif, closest_error, error_average, nbValues
 554        last_error, last_dif, closest_error, nbValues = 0.0, 0.0, 0.0, 0.0
 555
 556    def set_last_error(self, error, mass_spectrum_obj):
 557        """Set the last error.
 558
 559        Parameters
 560        ----------
 561        error : float
 562            The error.
 563        mass_spectrum_obj : MassSpectrum
 564            The mass spectrum object.
 565        """
 566        # set the changes to the global variables, not internal ones
 567        global last_error, last_dif, closest_error, error_average, nbValues
 568
 569        if mass_spectrum_obj.molecular_search_settings.error_method == "distance":
 570            dif = error - last_error
 571            if dif < last_dif:
 572                last_dif = dif
 573                closest_error = error
 574                mass_spectrum_obj.molecular_search_settings.min_ppm_error = (
 575                    closest_error
 576                    - mass_spectrum_obj.molecular_search_settings.mz_error_range
 577                )
 578                mass_spectrum_obj.molecular_search_settings.max_ppm_error = (
 579                    closest_error
 580                    + mass_spectrum_obj.molecular_search_settings.mz_error_range
 581                )
 582
 583        elif mass_spectrum_obj.molecular_search_settings.error_method == "lowest":
 584            if error < last_error:
 585                mass_spectrum_obj.molecular_search_settings.min_ppm_error = (
 586                    error - mass_spectrum_obj.molecular_search_settings.mz_error_range
 587                )
 588                mass_spectrum_obj.molecular_search_settings.max_ppm_error = (
 589                    error + mass_spectrum_obj.molecular_search_settings.mz_error_range
 590                )
 591                last_error = error
 592
 593        elif mass_spectrum_obj.molecular_search_settings.error_method == "symmetrical":
 594            mass_spectrum_obj.molecular_search_settings.min_ppm_error = (
 595                mass_spectrum_obj.molecular_search_settings.mz_error_average
 596                - mass_spectrum_obj.molecular_search_settings.mz_error_range
 597            )
 598            mass_spectrum_obj.molecular_search_settings.max_ppm_error = (
 599                mass_spectrum_obj.molecular_search_settings.mz_error_average
 600                + mass_spectrum_obj.molecular_search_settings.mz_error_range
 601            )
 602
 603        elif mass_spectrum_obj.molecular_search_settings.error_method == "average":
 604            nbValues += 1
 605            error_average = error_average + ((error - error_average) / nbValues)
 606            mass_spectrum_obj.molecular_search_settings.min_ppm_error = (
 607                error_average
 608                - mass_spectrum_obj.molecular_search_settings.mz_error_range
 609            )
 610            mass_spectrum_obj.molecular_search_settings.max_ppm_error = (
 611                error_average
 612                + mass_spectrum_obj.molecular_search_settings.mz_error_range
 613            )
 614
 615        else:
 616            # using set mass_spectrum_obj.molecular_search_settings.min_ppm_error  and max_ppm_error range
 617            pass
 618
 619        # returns the error based on the selected method at mass_spectrum_obj.molecular_search_settings.method
 620
 621    @staticmethod
 622    def calc_error(mz_exp, mz_calc, method="ppm"):
 623        """Calculate the error.
 624
 625        Parameters
 626        ----------
 627        mz_exp : float
 628            The experimental m/z value.
 629        mz_calc : float
 630            The calculated m/z value.
 631        method : str, optional
 632            The method, by default 'ppm'.
 633
 634        Raises
 635        -------
 636        Exception
 637            If the method is not ppm or ppb.
 638
 639        Returns
 640        -------
 641        float
 642            The error.
 643        """
 644
 645        if method == "ppm":
 646            multi_factor = 1_000_000
 647
 648        elif method == "ppb":
 649            multi_factor = 1_000_000_000
 650
 651        elif method == "perc":
 652            multi_factor = 100
 653
 654        else:
 655            raise Exception(
 656                "method needs to be ppm or ppb, you have entered %s" % method
 657            )
 658
 659        if mz_exp:
 660            return ((mz_exp - mz_calc) / mz_calc) * multi_factor
 661
 662        else:
 663            raise Exception("Please set mz_calc first")
 664
 665    def find_formulas(
 666        self,
 667        formulas,
 668        min_abundance,
 669        mass_spectrum_obj,
 670        ms_peak,
 671        ion_type,
 672        ion_charge,
 673        adduct_atom=None,
 674    ):
 675        """Find the formulas.
 676
 677        Parameters
 678        ----------
 679        formulas : list of MolecularFormula
 680            The list of molecular formulas.
 681        min_abundance : float
 682            The minimum abundance threshold.
 683        mass_spectrum_obj : MassSpectrum
 684            The mass spectrum object.
 685        ms_peak : MSPeak
 686            The mass spectrum peak.
 687        ion_type : str
 688            The ion type.
 689        ion_charge : int
 690            The ion charge.
 691        adduct_atom : str, optional
 692            The adduct atom, by default None.
 693
 694        Returns
 695        -------
 696        list of MSPeak
 697            The list of mass spectrum peaks with assigned molecular formulas.
 698
 699        Notes
 700        -----
 701        Uses the closest error the next search (this is not ideal, it needs to use confidence
 702        metric to choose the right candidate then propagate the error using the error from the best candidate).
 703        It needs to add s/n to the equation.
 704        It need optimization to define the mz_error_range within a m/z unit since it is directly proportional
 705        with the mass, and inversely proportional to the rp. It's not linear, i.e., sigma mass.
 706        The idea it to correlate sigma to resolving power, signal to noise and sample complexity per mz unit.
 707        Method='distance'
 708        """
 709        mspeak_assigned_index = list()
 710
 711        min_ppm_error = mass_spectrum_obj.molecular_search_settings.min_ppm_error
 712        max_ppm_error = mass_spectrum_obj.molecular_search_settings.max_ppm_error
 713
 714        min_abun_error = mass_spectrum_obj.molecular_search_settings.min_abun_error
 715        max_abun_error = mass_spectrum_obj.molecular_search_settings.max_abun_error
 716
 717        # f = open("abundance_error.txt", "a+")
 718        ms_peak_mz_exp, ms_peak_abundance = ms_peak.mz_exp, ms_peak.abundance
 719        # min_error = min([pmf.mz_error for pmf in possible_formulas])
 720
 721        def mass_by_ion_type(possible_formula_obj):
 722            if ion_type == Labels.protonated_de_ion:
 723                return possible_formula_obj._protonated_mz(ion_charge)
 724
 725            elif ion_type == Labels.radical_ion:
 726                return possible_formula_obj._radical_mz(ion_charge)
 727
 728            elif ion_type == Labels.adduct_ion and adduct_atom:
 729                return possible_formula_obj._adduct_mz(ion_charge, adduct_atom)
 730
 731            else:
 732                # will return externally calculated mz if is set, #use on Bruker Reference list import
 733                # if the ion type is known the ion mass based on molecular formula ion type
 734                # if ion type is unknow will return neutral mass
 735                return possible_formula_obj.mz_calc
 736
 737        if formulas:
 738            if isinstance(formulas[0], LCMSLibRefMolecularFormula):
 739                possible_mf_class = True
 740
 741            else:
 742                possible_mf_class = False
 743
 744        for possible_formula in formulas:
 745            if possible_formula:
 746                error = self.calc_error(
 747                    ms_peak_mz_exp, mass_by_ion_type(possible_formula)
 748                )
 749
 750                # error = possible_formula.mz_error
 751
 752                if min_ppm_error <= error <= max_ppm_error:
 753                    # update the error
 754
 755                    self.set_last_error(error, mass_spectrum_obj)
 756
 757                    # add molecular formula match to ms_peak
 758
 759                    # get molecular formula dict from sql obj
 760                    # formula_dict = pickle.loads(possible_formula.mol_formula)
 761                    # if possible_mf_class:
 762
 763                    #    molecular_formula = deepcopy(possible_formula)
 764
 765                    # else:
 766
 767                    formula_dict = possible_formula.to_dict()
 768                    # create the molecular formula obj to be stored
 769                    if possible_mf_class:
 770                        molecular_formula = LCMSLibRefMolecularFormula(
 771                            formula_dict,
 772                            ion_charge,
 773                            ion_type=ion_type,
 774                            adduct_atom=adduct_atom,
 775                        )
 776
 777                        molecular_formula.name = possible_formula.name
 778                        molecular_formula.kegg_id = possible_formula.kegg_id
 779                        molecular_formula.cas = possible_formula.cas
 780
 781                    else:
 782                        molecular_formula = MolecularFormula(
 783                            formula_dict,
 784                            ion_charge,
 785                            ion_type=ion_type,
 786                            adduct_atom=adduct_atom,
 787                        )
 788                    # add the molecular formula obj to the mspeak obj
 789                    # add the mspeak obj and it's index for tracking next assignment step
 790
 791                    if self.find_isotopologues:
 792                        # calculates isotopologues
 793                        isotopologues = molecular_formula.isotopologues(
 794                            min_abundance,
 795                            ms_peak_abundance,
 796                            mass_spectrum_obj.dynamic_range,
 797                        )
 798
 799                        # search for isotopologues
 800                        for isotopologue_formula in isotopologues:
 801                            molecular_formula.expected_isotopologues.append(
 802                                isotopologue_formula
 803                            )
 804                            # move this outside to improve preformace
 805                            # we need to increase the search space to -+1 m_z
 806                            first_index, last_index = (
 807                                mass_spectrum_obj.get_nominal_mz_first_last_indexes(
 808                                    isotopologue_formula.mz_nominal_calc
 809                                )
 810                            )
 811
 812                            for ms_peak_iso in mass_spectrum_obj[
 813                                first_index:last_index
 814                            ]:
 815                                error = self.calc_error(
 816                                    ms_peak_iso.mz_exp, isotopologue_formula.mz_calc
 817                                )
 818
 819                                if min_ppm_error <= error <= max_ppm_error:
 820                                    # need to define error distribution for abundance measurements
 821
 822                                    # if mass_spectrum_obj.is_centroid:
 823
 824                                    abundance_error = self.calc_error(
 825                                        isotopologue_formula.abundance_calc,
 826                                        ms_peak_iso.abundance,
 827                                        method="perc",
 828                                    )
 829
 830                                    # area_error = self.calc_error(ms_peak.area, ms_peak_iso.area, method='perc')
 831
 832                                    # margin of error was set empirically/ needs statistical calculation
 833                                    #  of margin of error for the measurement of the abundances
 834                                    if (
 835                                        min_abun_error
 836                                        <= abundance_error
 837                                        <= max_abun_error
 838                                    ):
 839                                        # update the error
 840
 841                                        self.set_last_error(error, mass_spectrum_obj)
 842
 843                                        # isotopologue_formula.mz_error = error
 844
 845                                        # isotopologue_formula.area_error = area_error
 846
 847                                        # isotopologue_formula.abundance_error = abundance_error
 848
 849                                        isotopologue_formula.mspeak_index_mono_isotopic = ms_peak.index
 850
 851                                        mono_isotopic_formula_index = len(ms_peak)
 852
 853                                        isotopologue_formula.mspeak_index_mono_isotopic = ms_peak.index
 854
 855                                        isotopologue_formula.mono_isotopic_formula_index = mono_isotopic_formula_index
 856
 857                                        # add mspeaks isotopologue index to the mono isotopic MolecularFormula obj and the respective formula position
 858
 859                                        # add molecular formula match to ms_peak
 860                                        x = ms_peak_iso.add_molecular_formula(
 861                                            isotopologue_formula
 862                                        )
 863
 864                                        molecular_formula.mspeak_mf_isotopologues_indexes.append(
 865                                            (ms_peak_iso.index, x)
 866                                        )
 867                                        # add mspeaks mono isotopic index to the isotopologue MolecularFormula obj
 868
 869                    y = ms_peak.add_molecular_formula(molecular_formula)
 870
 871                    mspeak_assigned_index.append((ms_peak.index, y))
 872
 873        return mspeak_assigned_index
 874
 875
 876class SearchMolecularFormulasLC:
 877    """Class for searching molecular formulas in a LC object.
 878
 879    Parameters
 880    ----------
 881    lcms_obj : LCMSBase
 882        The LCMSBase object.
 883    sql_db : MolForm_SQL, optional
 884        The SQL database object, by default None.
 885    first_hit : bool, optional
 886        Flag to indicate whether to skip peaks that already have a molecular formula assigned, by default False.
 887    find_isotopologues : bool, optional
 888        Flag to indicate whether to find isotopologues, by default True.
 889
 890    Methods
 891    -------
 892
 893    * search_spectra_against_candidates().
 894        Search a list of mass spectra against a list of candidate formulas with a given ion type and charge.
 895    * bulk_run_molecular_formula_search().
 896        Run the molecular formula search on the given list of mass spectra.
 897        Pulls the settings from the LCMSBase object to set ion type and charge to search for. 
 898    * run_mass_feature_search().
 899        Run the molecular formula search on mass features.
 900        Calls bulk_run_molecular_formula_search() with specified mass spectra and mass peaks.
 901    * run_untargeted_worker_ms1().
 902        Run untargeted molecular formula search on the ms1 mass spectrum.
 903        DEPRECATED: use run_mass_feature_search() or bulk_run_molecular_formula_search() instead.
 904    * run_target_worker_ms1().
 905        Run targeted molecular formula search on the ms1 mass spectrum.
 906        DEPRECATED: use run_mass_feature_search() or bulk_run_molecular_formula_search() instead.
 907    """
 908
 909    def __init__(self, lcms_obj, sql_db=None, first_hit=False, find_isotopologues=True):
 910        self.first_hit = first_hit
 911
 912        self.find_isotopologues = find_isotopologues
 913
 914        self.lcms_obj = lcms_obj
 915
 916        if not sql_db:
 917            self.sql_db = MolForm_SQL(
 918                url=self.lcms_obj.parameters.mass_spectrum['ms1'].molecular_search.url_database
 919            )
 920
 921        else:
 922            self.sql_db = sql_db
 923
 924    def search_spectra_against_candidates(self, mass_spectrum_list, ms_peaks_list, candidate_formulas, ion_type, ion_charge):
 925        """Search a list of mass spectra against a list of candidate formulas with a given ion type and charge.
 926
 927        Parameters
 928        ----------
 929        mass_spectrum_list : list of MassSpectrum
 930            The list of mass spectra to perform the search on.
 931        ms_peaks_list : list of lists of MSPeak objects
 932            The list of mass spectrum peaks to search within each mass spectrum.
 933        candidate_formulas : dict
 934            The candidate formulas.
 935        ion_type : str
 936            The ion type.
 937        ion_charge : int
 938            The ion charge, either 1 or -1.
 939
 940        Notes
 941        -----
 942        This function is designed to be used with the bulk_run_molecular_formula_search function.
 943        """
 944        for mass_spectrum, ms_peaks in zip(mass_spectrum_list, ms_peaks_list):
 945            single_ms_search = SearchMolecularFormulas(
 946                mass_spectrum,
 947                sql_db=self.sql_db,
 948                first_hit=self.first_hit,
 949                find_isotopologues=self.find_isotopologues,
 950            )
 951            single_ms_search.run_search(
 952                ms_peaks,
 953                candidate_formulas,
 954                mass_spectrum.min_abundance,
 955                ion_type,
 956                ion_charge,
 957            )
 958
 959    def bulk_run_molecular_formula_search(self, mass_spectrum_list, ms_peaks_list, mass_spectrum_setting_key='ms1'):
 960        """Run the molecular formula search on the given list of mass spectra
 961
 962        Parameters
 963        ----------
 964        mass_spectrum_list : list of MassSpectrum
 965            The list of mass spectra to search.
 966        ms_peaks_list : list of lists of MSPeak objects 
 967            The mass peaks to perform molecular formula search within each mass spectrum
 968        mass_spectrum_setting_key : str, optional
 969            The mass spectrum setting key, by default 'ms1'.
 970            This is used to get the appropriate molecular search settings from the LCMSBase object
 971        """
 972        # Set min_abundance and nominal_mzs
 973        if self.lcms_obj.polarity == "positive":
 974            ion_charge = 1
 975        elif self.lcms_obj.polarity == "negative":
 976            ion_charge = -1
 977        else:
 978            raise ValueError("Polarity must be either 'positive' or 'negative'")
 979        
 980        # Check that the length of the mass spectrum list and the ms_peaks list are the same
 981        if len(mass_spectrum_list) != len(ms_peaks_list):
 982            raise ValueError("The length of the mass spectrum list and the ms_peaks list must be the same")
 983        
 984        nominal_mzs = [x.nominal_mz for x in mass_spectrum_list]
 985        nominal_mzs = list(set([item for sublist in nominal_mzs for item in sublist]))
 986        verbose = self.lcms_obj.parameters.mass_spectrum[mass_spectrum_setting_key].molecular_search.verbose_processing 
 987
 988        # reset average error, only relevant if average mass error method is being used
 989        SearchMolecularFormulaWorker(
 990            find_isotopologues=self.find_isotopologues
 991        ).reset_error(mass_spectrum_list[0])
 992
 993        # check database for all possible molecular formula combinations based on the setting passed to self.mass_spectrum_obj.molecular_search_settings
 994        classes = MolecularCombinations(self.sql_db).runworker(
 995            self.lcms_obj.parameters.mass_spectrum[mass_spectrum_setting_key].molecular_search,
 996            print_time=self.lcms_obj.parameters.mass_spectrum[mass_spectrum_setting_key].molecular_search.verbose_processing
 997        )
 998        
 999        try:
1000            # split the database load to not blowout the memory
1001            for classe_chunk in chunks(
1002                classes, self.lcms_obj.parameters.mass_spectrum[mass_spectrum_setting_key].molecular_search.db_chunk_size
1003            ):
1004                classes_str_list = [class_tuple[0] for class_tuple in classe_chunk]
1005
1006                # load the molecular formula objs binned by ion type and heteroatoms classes, {ion type:{classe:[list_formula]}}
1007                # for adduct ion type a third key is added {atoms:{ion type:{classe:[list_formula]}}}
1008                dict_res = SearchMolecularFormulas.database_to_dict(
1009                    classes_str_list,
1010                    nominal_mzs,
1011                    self.lcms_obj.parameters.mass_spectrum[mass_spectrum_setting_key].molecular_search,
1012                    ion_charge,
1013                    sql_db=self.sql_db,
1014                )
1015
1016                pbar = tqdm.tqdm(classe_chunk, disable=not verbose)
1017                for classe_tuple in pbar:
1018                    # class string is a json serialized dict
1019                    classe_str = classe_tuple[0]
1020
1021                    # Perform search for (de)protonated ion type
1022                    if self.lcms_obj.parameters.mass_spectrum[mass_spectrum_setting_key].molecular_search.isProtonated:
1023                        ion_type = Labels.protonated_de_ion
1024
1025                        pbar.set_description_str(
1026                            desc="Started molecular formula search for class %s, (de)protonated "
1027                            % classe_str,
1028                            refresh=True,
1029                        )
1030
1031                        candidate_formulas = dict_res.get(ion_type).get(classe_str)
1032
1033                        if candidate_formulas:
1034                            self.search_spectra_against_candidates(
1035                                mass_spectrum_list=mass_spectrum_list,
1036                                ms_peaks_list=ms_peaks_list,
1037                                candidate_formulas=candidate_formulas,
1038                                ion_type=ion_type,
1039                                ion_charge=ion_charge,
1040                            )
1041
1042                    # Perform search for radical ion type
1043                    if self.lcms_obj.parameters.mass_spectrum[mass_spectrum_setting_key].molecular_search.isRadical:
1044                        pbar.set_description_str(
1045                            desc="Started molecular formula search for class %s, radical "
1046                            % classe_str,
1047                            refresh=True,
1048                        )
1049
1050                        ion_type = Labels.radical_ion
1051
1052                        candidate_formulas = dict_res.get(ion_type).get(classe_str)
1053
1054                        if candidate_formulas:
1055                            self.search_spectra_against_candidates(
1056                                mass_spectrum_list=mass_spectrum_list,
1057                                ms_peaks_list=ms_peaks_list,
1058                                candidate_formulas=candidate_formulas,
1059                                ion_type=ion_type,
1060                                ion_charge=ion_charge,
1061                            )
1062
1063                    # Perform search for adduct ion type
1064                    # looks for adduct, used_atom_valences should be 0
1065                    # this code does not support H exchance by halogen atoms
1066                    if self.lcms_obj.parameters.mass_spectrum[mass_spectrum_setting_key].molecular_search.isAdduct:
1067                        pbar.set_description_str(
1068                            desc="Started molecular formula search for class %s, adduct "
1069                            % classe_str,
1070                            refresh=True,
1071                        )
1072
1073                        ion_type = Labels.adduct_ion
1074                        dict_atoms_formulas = dict_res.get(ion_type)
1075
1076                        for adduct_atom, dict_by_class in dict_atoms_formulas.items():
1077                            candidate_formulas = dict_by_class.get(classe_str)
1078
1079                            if candidate_formulas:
1080                                self.search_spectra_against_candidates(
1081                                    mass_spectrum_list=mass_spectrum_list,
1082                                    ms_peaks_list=ms_peaks_list,
1083                                    candidate_formulas=candidate_formulas,
1084                                    ion_type=ion_type,
1085                                    ion_charge=ion_charge,
1086                                )
1087        finally:
1088            self.sql_db.close()
1089        
1090    def run_mass_feature_search(self):
1091        """Run the molecular formula search on the mass features.
1092        
1093        Calls bulk_run_molecular_formula_search() with specified mass spectra and mass peaks.
1094        """
1095        mass_features_df = self.lcms_obj.mass_features_to_df()
1096
1097        # Get the list of mass spectrum (and peaks to search with each mass spectrum) for all mass features
1098        scan_list = mass_features_df.apex_scan.unique()
1099        mass_spectrum_list = [self.lcms_obj._ms[x] for x in scan_list]
1100        ms_peaks = []
1101        for scan in scan_list:
1102            mf_df_scan = mass_features_df[mass_features_df.apex_scan == scan]
1103            peaks_to_search = [
1104                self.lcms_obj.mass_features[x].ms1_peak for x in mf_df_scan.index.tolist()
1105            ]
1106            ms_peaks.append(peaks_to_search)
1107        
1108        # Run the molecular formula search
1109        self.bulk_run_molecular_formula_search(mass_spectrum_list, ms_peaks)
1110    
1111    def run_untargeted_worker_ms1(self):
1112        """Run untargeted molecular formula search on the ms1 mass spectrum."""
1113        raise NotImplementedError("run_untargeted_worker_ms1 search is not implemented in CoreMS 3.0 and greater")
1114
1115    def run_target_worker_ms1(self):
1116        """Run targeted molecular formula search on the ms1 mass spectrum."""
1117        raise NotImplementedError("run_target_worker_ms1 formula search is not yet implemented in CoreMS 3.0 and greater")