def simulation(arguments: List[str] = sys.argv[1:]):
args = _arguments(arguments)
config = dict()
if args.type == "config":
if isinstance(args.config, Path) or args.example:
config = _handle_configuration(
args) # Will exit if printing config
else:
print("Try 'pysimulate config --help")
sys.exit(1)
print("Parsing provided parameters")
function_path, intensity_name, setup_name, \
parameters, data_path, slot_name, output, use_cache = \
_process_arguments(config, args)
print("Loading data")
if slot_name:
factory = slot_table.SlotFactory(data_path, "r")
data = factory.get_slot(slot_name)
else:
data = processor.DataProcessor(use_cache).parse(data_path)
print("Loading functions")
try:
intensity = function.load(function_path, intensity_name)
except ImportError as error:
print("Failed to load the intensity function!")
raise error
try:
setup = function.load(function_path, setup_name)
except ImportError as error:
print("Failed to load the setup function!")
raise error
print("Starting Simulation")
rejection = simulate.calculate_intensities(intensity, setup, parameters,
data)
if slot_name:
data.add_data("rejection", rejection)
if output:
processor.DataProcessor().write(output, rejection)
if not slot_name and not output:
_LOGGER.warn("No output was provided! Using stdout for output!")
[print(int(x)) for x in rejection]
def _table_to_directory(folder, base_dir):
parser = processor.DataProcessor()
folder_pbar = tqdm.tqdm(folder.binning.get_bin_directory(), unit="folder")
file_pbar = tqdm.tqdm(total=len(folder.unmanaged.nodes) + 1, unit="file")
folder_pbar.set_description()
for directory in folder_pbar:
# Create the new directory for the bin
file_pbar.reset()
new_path = base_dir / directory.bin_location
new_path.mkdir(parents=True)
root_path = new_path / directory.root.filename
with parser.get_writer(root_path, directory.root.data_type) as w:
for event in tqdm.tqdm(directory.root.iterate_data(1),
total=len(directory),
unit=" event",
postfix=f"file={folder.root.filename}",
position=2):
w.write(event)
file_pbar.update(1)
for extra_data in directory.unmanaged:
data_path = new_path / extra_data.filename
if extra_data.data_type == processor.DataType.STRUCTURED:
with parser.get_writer(data_path, extra_data.data_type) as w:
for event in tqdm.tqdm(
extra_data.iterate_data(1),
unit="event",
total=len(extra_data),
postfix=f"file={folder.root.filename}",
position=2):
w.write(event)
else:
parser.write(data_path, extra_data.data)
file_pbar.update(1)
folder_pbar.close()
file_pbar.close()
def parser():
return processor.DataProcessor(False, False)
def start_masking(arguments: List[str] = sys.argv[1:]):
args = _arguments(arguments)
file_manager = processor.DataProcessor(False)
# Load the input file
if args.input.exists():
input_file = file_manager.get_reader(args.input)
output_file = file_manager.get_writer(args.output, input_file.data_type)
else:
print(f"{args.input} must exist!")
return 1
# Load the correct file
if args.use_or:
logic = npy.logical_or
elif args.use_xor:
logic = npy.logical_xor
elif args.use_or and args.use_xor:
print("Only select OR or XOR, not both!")
return 1
else:
logic = npy.logical_and
# Setup progress bar for mask files
if len(args.mask) > 1:
disable = False
else:
disable = True
# Merge together masks
pf = None
for mask_file in tqdm.tqdm(args.mask, disable=disable): # type: Path
if mask_file.exists():
current_pf = file_manager.parse(mask_file)
else:
print(f"{mask_file} must exist!")
return 1
# Convert selection array to a boolean mask.
if "u8" == current_pf.dtype or "u4" == current_pf.dtype:
new_pf = npy.zeros(len(input_file), bool)
new_pf[current_pf] = True
current_pf = new_pf
if isinstance(pf, type(None)):
pf = current_pf
else:
pf = logic(pf, current_pf)
# Handle no masks provided
if isinstance(pf, type(None)):
pf = npy.ones(len(input_file), bool)
if len(pf) != len(input_file):
print(
f"Masking data isn't the same length as input!"
f" Mask is {len(pf)} and input is {len(input_file)}."
)
return 1
# Setup description
if pf.all() == 1:
description = "Converting:"
else:
description = "Masking:"
# Input masked to output
progress = tqdm.tqdm(
zip(pf, input_file), total=len(pf), unit="Events",
desc=description
)
for do_write, event in progress:
if do_write:
output_file.write(event)
input_file.close()
output_file.close()
def start_binning(arguments: List[str] = sys.argv[1:]):
table = Path("binning.hd5")
parser = processor.DataProcessor()
base_dir = Path("bin_data")
args = _arguments(arguments)
if args.example:
print(yaml.dump(CONFIGURATION))
sys.exit(0)
if not args.configuration:
print("You must provide a configuration file!")
sys.exit(1)
with open(args.configuration) as stream:
config = yaml.load(stream, Loader=yaml.FullLoader)
if base_dir.exists():
print(f"Please delete {base_dir} first before running.")
sys.exit(1)
if table.exists():
table.unlink()
root_data = parser.get_reader(config["primary data"])
# Calculate total space requirements
event_count = root_data.get_event_count()
num_bins = 1
for value in config["binning settings"].values():
if value["type"] == "range":
num_bins *= value["count"] + 2
else:
num_bins += value["count"]
space = (num_bins * root_data.get_event_count())
space += root_data.input_path.stat().st_size
if config["binned data"]:
for file in config["binned data"]:
space += Path(file).stat().st_size
space /= 1000000000 # Bytes to Gigabytes
print(f"Estimated to take {space:.2f}Gb of disk space")
folder_manager = project.ProjectDatabase(table, "w")
folder = folder_manager.make_folder("binning_data", root_data,
config["primary data"], True)
if config["binned data"]:
pbar = tqdm.tqdm(total=len(config["binned data"]),
unit="file",
desc="Parsing binned data")
for file in config["binned data"]:
pbar.set_postfix_str(f"file={file}")
data = parser.parse(file)
folder.unmanaged.add(file.replace(".", "_"), data, file)
pbar.update(1)
pbar.close()
for setting_index in config["binning settings"].keys():
settings = config["binning settings"][setting_index]
# Convert configuration to variable
if settings["variable"] == "mass":
var = folder.binning.BinVars.MASS
elif settings["variable"] == "beam":
var = folder.binning.BinVars.BEAM
elif settings["variable"] == "t":
var = folder.binning.BinVars.T
elif settings["variable"] == "tprime":
var = folder.binning.BinVars.T_PRIME
else:
print(f"{settings['variable']} variable not understood in "
f"{setting_index}!")
sys.exit(1)
# Queue the binning parameters
if settings["type"] == "range":
folder.binning.add_fixed_range(var, settings["lower limit"],
settings["upper limit"],
settings["count"])
elif settings["type"] == "fixed":
folder.binning.add_fixed_count(var, settings["count"])
else:
print(f"{settings['type']} is not understood in "
f"{setting_index}!")
sys.exit(1)
folder.binning.execute()
_table_to_directory(folder, base_dir)
folder_manager.close()
if "clean up" in config:
if config["clean up"]:
table.unlink()
else:
table.unlink()
print("\nCompleted!")