diff --git a/backlogger/find.py b/backlogger/find.py new file mode 100644 index 0000000..b3f7996 --- /dev/null +++ b/backlogger/find.py @@ -0,0 +1,91 @@ +import sqlite3 +import datalad.api as dl +import os +import json +import pandas as pd +# this will implement the search functionality + + +def _project_aka_lookup(alias): + # this will lookup the project name based on the alias + conn = sqlite3.connect('backlogger.db') + c = conn.cursor() + c.execute(f"SELECT * FROM 'projects' WHERE alias = '{alias}'") + results = c.fetchall() + conn.close() + if len(results) > 1: + print("Error: multiple projects found with alias " + alias) + elif len(results) == 0: + raise Exception("Error: no project found with alias " + alias) + return results[0][0] + + +def _db_lookup(db, ensemble, correlator_name, project=None, code=None, parameters=None, created_before=None, created_after=None, updated_before=None, updated_after=None, revision=None): + project_str = project + + search_expr = f"SELECT * FROM 'backlogs' WHERE name = '{correlator_name}' AND ensemble = '{ensemble}'" + if project: + search_expr += f" AND project = '{project_str}'" + if code: + search_expr += f" AND code = '{code}'" + if parameters: + search_expr += f" AND parameters = '{parameters}'" + if created_before: + search_expr += f" AND created_at < '{created_before}'" + if created_after: + search_expr += f" AND created_at > '{created_after}'" + if updated_before: + search_expr += f" AND updated_at < '{updated_before}'" + if updated_after: + search_expr += f" AND updated_at > '{updated_after}'" + print(search_expr) + conn = sqlite3.connect(db) + results = pd.read_sql(search_expr, conn) + conn.close() + return results + + +def filter_results(results, **kwargs): + drops = [] + for ind in range(len(results)): + result = results.iloc[ind] + if result['code'] == 'sfcf': + param = json.loads(result['parameters']) + if 'offset' in kwargs: + print("checking offset") + if kwargs.get('offset') != param['offset']: + print("dropping") + drops.append(ind) + continue + if 'quark_masses' in kwargs: + quark_masses = kwargs['quark_masses'] + if (quark_masses[0] != param['quarks'][0]['mass'] or quark_masses[1] != param['quarks'][1]['mass']) and (quark_masses[0] != param['quarks'][1]['mass'] or quark_masses[1] != param['quarks'][0]['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']): + drops.append(ind) + continue + if 'wf1' in kwargs: + wf1 = kwargs['wf1'] + if (wf1[0] != param['wf1'][0]) or (wf1[1][0] != param['wf1'][1][0]) or (wf1[1][1] != param['wf1'][1][1]): + drops.append(ind) + continue + if 'wf2' in kwargs: + wf2 = kwargs['wf2'] + if (wf2[0] != param['wf2'][0]) or (wf2[1][0] != param['wf2'][1][0]) or (wf2[1][1] != param['wf2'][1][1]): + drops.append(ind) + continue + + return results.drop(drops) + + +def find_correlator(path, ensemble, correlator_name, project=None, code=None, parameters=None, created_before=None, created_after=None, updated_before=None, updated_after=None, revision=None, **kwargs): + db = path + '/backlogger.db' + if os.path.exists(db): + dl.get(db, dataset=path) + results = _db_lookup(db, ensemble, correlator_name, project, code=None, parameters=None, created_before=created_before, created_after=created_after, updated_before=updated_before, updated_after=updated_after, revision=revision) + results = filter_results(results, **kwargs) + print("Found " + str(len(results)) + " results") + return results