From 4f7d9d71b35074e39c6666ed6950b70122c2afe4 Mon Sep 17 00:00:00 2001
From: Justus Kuhlmann <j_kuhl19@uni-muenster.de>
Date: Tue, 13 May 2025 09:45:37 +0000
Subject: [PATCH] Implement better finding for sfcf projects

---
 corrlib/find.py  | 29 ++++++++++++++++++++++++++---
 corrlib/tools.py |  7 +++++++
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/corrlib/find.py b/corrlib/find.py
index a1b25f9..432bb26 100644
--- a/corrlib/find.py
+++ b/corrlib/find.py
@@ -5,6 +5,7 @@ import json
 import pandas as pd
 import numpy as np
 from .input.implementations import codes
+from .tools import m2k, k2m
 # this will implement the search functionality
 
 
@@ -65,11 +66,33 @@ def sfcf_filter(results, **kwargs):
                 if kwargs.get('offset') != param['offset']:
                     drops.append(ind)
                     continue
-            if 'quark_masses' in kwargs:
-                quark_masses = kwargs['quark_masses']
-                if (not np.isclose(quark_masses[0], param['quarks'][0]['mass']) or not np.isclose(quark_masses[1], param['quarks'][1]['mass'])):
+            if 'quark_kappas' in kwargs:
+                kappas = kwargs['quark_kappas']
+                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 'qm1' in kwargs:
+                quark_mass1 = kwargs['qm1']
+                if not isinstance(quark_mass1, list):
+                    if (not np.isclose(quark_mass1, k2m(param['quarks'][0]['mass']))):
+                        drops.append(ind)
+                        continue
+                else:
+                    if len(quark_mass1) == 2:
+                        if (quark_mass1[0] > k2m(param['quarks'][0]['mass'])) or (quark_mass1[1] < k2m(param['quarks'][0]['mass'])):
+                            drops.append(ind)
+                            continue
+            if 'qm2' in kwargs:
+                quark_mass2 = kwargs['qm2']
+                if not isinstance(quark_mass2, list):
+                    if (not np.isclose(quark_mass2, k2m(param['quarks'][1]['mass']))):
+                        drops.append(ind)
+                        continue
+                else:
+                    if len(quark_mass2) == 2:
+                        if (quark_mass2[0] > k2m(param['quarks'][1]['mass'])) or (quark_mass2[1] < k2m(param['quarks'][1]['mass'])):
+                            drops.append(ind)
+                            continue
             if 'quark_thetas' in kwargs:
                 quark_thetas = kwargs['quark_thetas']
                 if (quark_thetas[0] != param['quarks'][0]['thetas'] and quark_thetas[1] != param['quarks'][1]['thetas']) or (quark_thetas[0] != param['quarks'][1]['thetas'] and quark_thetas[1] != param['quarks'][0]['thetas']):
diff --git a/corrlib/tools.py b/corrlib/tools.py
index e424b59..3ac8bfe 100644
--- a/corrlib/tools.py
+++ b/corrlib/tools.py
@@ -9,3 +9,10 @@ def list2str(mylist):
     return s
 
 cached = True
+
+def m2k(m):
+    return 1/(2*m+8)
+
+
+def k2m(k):
+    return (1/(2*k))-4