diff --git a/corrlib/find.py b/corrlib/find.py index 432bb26..0c5f9a6 100644 --- a/corrlib/find.py +++ b/corrlib/find.py @@ -5,7 +5,7 @@ import json import pandas as pd import numpy as np from .input.implementations import codes -from .tools import m2k, k2m +from .tools import k2m # this will implement the search functionality @@ -71,6 +71,33 @@ def sfcf_filter(results, **kwargs): if (not np.isclose(kappas[0], param['quarks'][0]['mass']) or not np.isclose(kappas[1], param['quarks'][1]['mass'])): drops.append(ind) continue + if 'quark_masses' in kwargs: + masses = kwargs['quark_masses'] + if (not np.isclose(masses[0], k2m(param['quarks'][0]['mass'])) or not np.isclose(masses[1], k2m(param['quarks'][1]['mass']))): + drops.append(ind) + continue + if 'qk1' in kwargs: + quark_kappa1 = kwargs['qk1'] + if not isinstance(quark_kappa1, list): + if (not np.isclose(quark_kappa1, param['quarks'][0]['mass'])): + drops.append(ind) + continue + else: + if len(quark_kappa1) == 2: + if (quark_kappa1[0] > param['quarks'][0]['mass']) or (quark_kappa1[1] < param['quarks'][0]['mass']): + drops.append(ind) + continue + if 'qk2' in kwargs: + quark_kappa2 = kwargs['qk2'] + if not isinstance(quark_kappa2, list): + if (not np.isclose(quark_kappa2, param['quarks'][1]['mass'])): + drops.append(ind) + continue + else: + if len(quark_kappa2) == 2: + if (quark_kappa2[0] > param['quarks'][1]['mass']) or (quark_kappa2[1] < param['quarks'][1]['mass']): + drops.append(ind) + continue if 'qm1' in kwargs: quark_mass1 = kwargs['qm1'] if not isinstance(quark_mass1, list): @@ -109,7 +136,6 @@ def sfcf_filter(results, **kwargs): if not (np.isclose(wf2[0][0], param['wf2'][0][0], 1e-8) and np.isclose(wf2[0][1][0], param['wf2'][0][1][0], 1e-8) and np.isclose(wf2[0][1][1], param['wf2'][0][1][1], 1e-8)): drops.append(ind) continue - return results.drop(drops)