bic/scripts/extract_objective.py

0001 """Extract objective from a file
0002
0003 Defines small function to extract
0004 a specified value from a JSON file.
0005
0006 Project: AID2E v0.0.0 - AI assisted Detector Design for EIC
0007 """
0008 import json
0009
0010 def ExtractObjective(*, extra_args, **_):
0011     """Extract objective
0012
0013        Extracts a metric based from the JSON file
0014        written by CalculateHitAngReso. The key
0015        of the value to extract and the path to
0016        the JSON file must be provided through
0017        the `extra_args` block.
0018
0019        Args:
0020            extra_args: dictionary of extra arugments
0021                        must contain `key` and `file`
0022        Returns:
0023            extracted metric formatted as a dictionary
0024     """
0025     if 'key' not in extra_args:
0026         raise KeyError(f"'key' not found in extra_args! extra_args = {extra_args}")
0027
0028     if 'file' not in extra_args:
0029         raise KeyError(f"'file' not found in extra_args! extra_args = {extra_args}")
0030     key  = extra_args['key']
0031     file = extra_args['file']
0032
0033     # relevant file has the same name as the ROOT
0034     # output but w/ the .root suffix changed to
0035     # .json
0036     #   - FIXME this is clunky! Maybe this is a
0037     #     a good use case for the artifact
0038     #     spec?
0039     file = file.replace('.root', '.json')
0040
0041     metrics = {}
0042     with open(file, 'r') as f:
0043         data         = json.load(f)
0044         value        = data[key]
0045         metrics[key] = value
0046     return metrics