diff --git a/dap/algos/__init__.py b/dap/algos/__init__.py index d5e8f0b..b70958a 100644 --- a/dap/algos/__init__.py +++ b/dap/algos/__init__.py @@ -2,6 +2,7 @@ from .addmask import calc_apply_additional_mask from .addmaskfile import calc_apply_additional_mask_from_file from .aggregation import calc_apply_aggregation +from .custom import calc_custom from .jfdata import JFData from .mask import calc_mask_pixels from .peakfind import calc_peakfinder_analysis diff --git a/dap/algos/custom.py b/dap/algos/custom.py new file mode 100644 index 0000000..214c7eb --- /dev/null +++ b/dap/algos/custom.py @@ -0,0 +1,64 @@ +import functools +import importlib.util as ilu +from pathlib import Path + + +BASE = "/gpfs/photonics/swissfel/buffer/dap/custom_dap_scripts" +BASE = Path(BASE) + + +def calc_custom(results, image, pixel_mask): + script = results.get("custom_script") + if not script: + return + + try: + func = load_custom(script) + res = func(results, image, pixel_mask) + except Exception as e: + en = type(e).__name__ + error = f"{en}: {e}" + return {"custom_script_error": error} + + if isinstance(res, dict): + res = {f"{script}:{k}": v for k, v in res.items()} + else: + res = {script: res} + + results.update(res) + + +@functools.cache +def load_custom(script): + beamline, name = script.split(":") + + ext = ".py" + if not name.endswith(ext): + name += ext + + fn = BASE / beamline / name + return load_proc(fn) + + +def load_proc(fn): + mod = load_module(fn) + + proc_func_name = "proc" + mod_name = mod.__name__ + + func = getattr(mod, proc_func_name, None) or getattr(mod, mod_name, None) + if func is None: + raise AttributeError(f'module "{mod_name}" contains neither "{proc_func_name}" nor "{mod_name}" function') + + return func + + +def load_module(file_path, module_name=None): + module_name = module_name or Path(file_path).stem + spec = ilu.spec_from_file_location(module_name, file_path) + module = ilu.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + + diff --git a/dap/worker.py b/dap/worker.py index 9df29c8..2b96866 100644 --- a/dap/worker.py +++ b/dap/worker.py @@ -5,7 +5,7 @@ import numpy as np from logzero import logger as log from algos import ( - calc_apply_aggregation, calc_apply_threshold, calc_mask_pixels, calc_peakfinder_analysis, + calc_apply_aggregation, calc_apply_threshold, calc_custom, calc_mask_pixels, calc_peakfinder_analysis, calc_radial_integration, calc_roi, calc_spi_analysis, calc_streakfinder_analysis, JFData ) from utils import Aggregator, BufferedJSON, make_bsread_timestamp, randskip, read_bit @@ -123,6 +123,7 @@ def work(backend_addr, accumulator_addr, visualisation_addr, fn_config, skip_fra calc_roi(results, pfimage) calc_spi_analysis(results, pfimage) calc_peakfinder_analysis(results, pfimage, pixel_mask) + calc_custom(results, pfimage, pixel_mask) # streak finder for convergent-beam diffraction experiments # changes image and pixel_mask in place if do_snr=True in parameters file