jkuhl/corrlib
1
0
Fork 0
Code Issues 11 Pull requests Releases Packages Activity Actions

add customtFilter

Browse source
This commit is contained in:
Justus Kuhlmann 2026-03-24 09:24:12 +01:00
commit 3fd557f3ee
Signed by: jkuhl
GPG key ID: 00ED992DD79B85A6
1 changed files with 15 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

20
corrlib/find.py
View file

@ -6,9 +6,10 @@ import numpy as np
from .input.implementations import codes from .input.implementations import codes
from .tools import k2m, get_db_file from .tools import k2m, get_db_file
from .tracker import get from .tracker import get
from typing import Any, Optional from typing import Any, Optional, Union
from pathlib import Path from pathlib import Path
import datetime as dt import datetime as dt
from collections.abc import Callable
def _project_lookup_by_alias(db: Path, alias: str) -> str: def _project_lookup_by_alias(db: Path, alias: str) -> str:
@ -251,22 +252,31 @@ def sfcf_filter(results: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
return results.drop(drops) return results.drop(drops)
def openQCD_filter(results:pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
return results
def find_record(path: Path, ensemble: str, correlator_name: str, code: str, project: Optional[str]=None, parameters: Optional[str]=None, def find_record(path: Path, ensemble: str, correlator_name: str, code: str, project: Optional[str]=None, parameters: Optional[str]=None,
created_before: Optional[str]=None, created_after: Optional[str]=None, updated_before: Optional[str]=None, updated_after: Optional[str]=None, revision: Optional[str]=None, **kwargs: Any) -> pd.DataFrame: created_before: Optional[str]=None, created_after: Optional[str]=None, updated_before: Optional[str]=None, updated_after: Optional[str]=None,
revision: Optional[str]=None,
customFilter: Optional[Callable[[pd.DataFrame], pd.DataFrame]] = None,
**kwargs: Any) -> pd.DataFrame:
db_file = get_db_file(path) db_file = get_db_file(path)
db = path / db_file db = path / db_file
if code not in codes: if code not in codes:
raise ValueError("Code " + code + "unknown, take one of the following:" + ", ".join(codes)) raise ValueError("Code " + code + "unknown, take one of the following:" + ", ".join(codes))
get(path, db_file) get(path, db_file)
results = _db_lookup(db, ensemble, correlator_name,code, project, parameters=parameters) results = _db_lookup(db, ensemble, correlator_name,code, project, parameters=parameters)
if any(arg is not None for arg in [created_before, created_after, updated_before, updated_after]): if any([arg is not None for arg in [created_before, created_after, updated_before, updated_after]]):
results = _time_filter(results, created_before, created_after, updated_before, updated_after) results = _time_filter(results, created_before, created_after, updated_before, updated_after)
if customFilter is not None:
results = customFilter(results)
if code == "sfcf": if code == "sfcf":
results = sfcf_filter(results, **kwargs) results = sfcf_filter(results, **kwargs)
elif code == "openQCD": elif code == "openQCD":
pass results = openQCD_filter(results, **kwargs)
else: else:
raise Exception raise ValueError(f"Code {code} is not known.")
print("Found " + str(len(results)) + " result" + ("s" if len(results)>1 else "")) print("Found " + str(len(results)) + " result" + ("s" if len(results)>1 else ""))
return results.reset_index() return results.reset_index()