diff --git a/benchmarks/imaging_shower_ML/run_benchmark.py b/benchmarks/imaging_shower_ML/run_benchmark.py index bc91124a8af24c3b91beb7fb198ed767a460df80..a196fcbecf83d6293b98ea52d6fdadc0e370ea3c 100755 --- a/benchmarks/imaging_shower_ML/run_benchmark.py +++ b/benchmarks/imaging_shower_ML/run_benchmark.py @@ -26,14 +26,14 @@ FILE_NAMES = dict( prep_script = os.path.join(SDIR, 'scripts', 'ml_data_preprocess.py'), ml_script = os.path.join(SDIR, 'scripts', 'ml_training.py'), - sim_dir = os.path.join('{outdir}', 'sim_data'), - epscan_dir = os.path.join('{outdir}', 'epcut'), - ml_dir = os.path.join('{outdir}', 'ml_result'), + sim_dir = os.path.join('{outdir}', '{run_type}', 'sim_data'), + epscan_dir = os.path.join('{outdir}', '{run_type}', 'epcut'), + ml_dir = os.path.join('{outdir}', '{run_type}', 'ml_result'), - gen_file = os.path.join('{outdir}', 'sim_data', '{ntag}_gen.hepmc'), - sim_file = os.path.join('{outdir}', 'sim_data', '{ntag}_sim.edm4hep.root'), - rec_file = os.path.join('{outdir}', 'sim_data', '{ntag}_rec.root'), - ml_file = os.path.join('{outdir}', 'sim_data', '{ntag}_data.h5'), + gen_file = os.path.join('{outdir}', '{run_type}', 'sim_data', '{ntag}_gen.hepmc'), + sim_file = os.path.join('{outdir}', '{run_type}', 'sim_data', '{ntag}_sim.edm4hep.root'), + rec_file = os.path.join('{outdir}', '{run_type}', 'sim_data', '{ntag}_rec.root'), + ml_file = os.path.join('{outdir}', '{run_type}', 'sim_data', '{ntag}_data.h5'), ) # default values for argument parser DEFAULT_COMPACT = os.path.join( @@ -124,11 +124,17 @@ if __name__ == '__main__': default='sim_output', help='output directory.' ) + parser.add_argument( + '-r', '--run-type', type=str, + dest='run_type', + default='epic_imcal', + help='a name specify the run type.' + ) parser.add_argument( '-t', '--name-tag', type=str, dest='ntag', default='imcal_ml', - help='name tag for output files.' + help='a name tag for output files.' ) parser.add_argument( '-c', '--compact', type=str, @@ -271,11 +277,13 @@ if __name__ == '__main__': # save run information # NOTE: runs with the same outdir will be combined (overwrite if nametag is the same) - run_data = {args.ntag: kwargs} + run_data = {args.run_type: {args.ntag: kwargs}} try: with open(os.path.join(args.outdir, 'result.json'), 'r') as f: run_data = json.load(f) - run_data.update({args.ntag: kwargs}) + run_info = run_data.get(args.run_type, {}) + run_info.update({args.ntag: kwargs}) + run_data[args.run_type] = run_info except (FileNotFoundError, json.decoder.JSONDecodeError): pass