diff --git a/src/kbmod_wf/task_impls/kbmod_search.py b/src/kbmod_wf/task_impls/kbmod_search.py index bace169..1ed8496 100644 --- a/src/kbmod_wf/task_impls/kbmod_search.py +++ b/src/kbmod_wf/task_impls/kbmod_search.py @@ -1,7 +1,11 @@ +import kbmod +from kbmod.work_unit import WorkUnit + +import os import time -def kbmod_search(input_wu=None, result_file=None, logger=None): +def placeholder_kbmod_search(input_wu=None, result_file=None, logger=None): logger.info("In the kbmod_search task_impl") with open(input_wu, "r") as f: for line in f: @@ -14,3 +18,59 @@ def kbmod_search(input_wu=None, result_file=None, logger=None): f.write(f"Logged: {value} - {time.time()}\n") return result_file + + +def kbmod_search( + input_wu=None, search_config_filepath=None, runtime_config={}, result_file=None, logger=None +): + kbmod_searcher = KBMODSearcher( + input_wu_filepath=input_wu, + search_config_filepath=search_config_filepath, + runtime_config=runtime_config, + output_filepath=result_file, + logger=logger, + ) + + return kbmod_searcher.run_search() + + +class KBMODSearcher: + def __init__(self, input_wu_filepath, search_config_filepath, runtime_config, output_filepath, logger): + self.input_wu_filepath = input_wu_filepath + self.search_config_filepath = search_config_filepath + self.runtime_config = runtime_config + self.output_filepath = output_filepath + self.logger = logger + + def run_search(self): + self.logger.info("Loading workunit from file") + wu = WorkUnit.from_fits(self.input_wu_filepath) + + self.logger.debug("Loaded work unit") + if self.search_config is not None: + # Load a search configuration, otherwise use the one loaded with the work unit + wu.config = kbmod.configuration.SearchConfiguration.from_file(self.search_config) + + config = wu.config + + # Modify the work unit results to be what is specified in command line args + input_parameters = { + "res_filepath": args.result_dir, + "result_filename": os.path.join(args.result_dir, "full_results.ecsv"), + } + config.set_multiple(input_parameters) + + # Save the search config in the results directory for record keeping + config.to_file(os.path.join(args.result_dir, "search_config.yaml")) + wu.config = config + + self.logger.info("Running KBMOD search") + res = kbmod.run_search.SearchRunner().run_search_from_work_unit(wu) + + self.logger.info("Search complete") + self.logger.info(f"Search results: {res}") + + self.logger.info("writing results table") + res.write_table(self.output_filepath) + + return self.output_filepath