diff --git a/python/_util.py b/python/_util.py index 40d4db57..131ef679 100644 --- a/python/_util.py +++ b/python/_util.py @@ -464,6 +464,8 @@ def _find_model_folder_and_file(base_dir, model_name, must_exist, specific_file= specific_file_path = os.path.join(base_dir, name, specific_file) if os.path.isfile(specific_file_path): return name, True, specific_file_path + else: + return name, True, None else: return name, True, None else: @@ -557,6 +559,108 @@ def _get_model_path(model_name, prefix=None, suffix=None, is_file=True, must_exi return os.path.join(base_dir, model_name, model_file_name) +def _get_model_folder(base_dir, model_name, must_exist): + model_names = [ + f for f in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, f)) + ] + + exact_model_names = [f for f in model_names if f.lower() == model_name.lower()] + + if len(exact_model_names) == 0: + model_names = [f for f in model_names if f.lower().startswith(model_name.lower())] + else: + model_names = exact_model_names + + if len(model_names) == 0 and must_exist: + raise ValueError(f"No model folders found for {model_name}\nSearched in {base_dir}") + elif len(model_names) == 0 and not must_exist: + return model_name, False + elif len(model_names) == 1: + return model_names[0], True + else: + raise ValueError(f"Multiple directories found for {model_name}\nSearched in {base_dir}") + +def _get_model_subfolders(base_dir, model_regex): + model_subfolders = [ + f for f in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, f)) + ] + model_subfolders = [ + f for f in model_subfolders if model_regex.match(f) is not None + ] + return model_subfolders + + +def _get_model_path(model_name, prefix=None, suffix=None, is_file=True, must_exist=True, specific_file=None): + _model_regex = re.compile( + r"^\s*" + _MODEL_PATTERN + ("" if suffix is None else r"(?:" + suffix + r")?") + r"\s*$", + re.VERBOSE | re.IGNORECASE, + ) + suffix = "" if suffix is None else suffix + + resources_dir = resource_package_dir() + base_dir = _get_base_directory(resources_dir, prefix) + + d = _model_regex.match(model_name) + if d is None: + raise ValueError(f"Invalid model name: {model_name}") + d = d.groupdict() + model_search_name, version = d["model_name"], d["version"] + + if version is not None: + version = normalize_version(version) + + found_model_name, folder_exists = _get_model_folder(base_dir, model_search_name, must_exist) + + model_dir = os.path.join(base_dir, found_model_name) + + if not must_exist and not folder_exists: + if version is None: + version = "v1" + + model_dir = os.path.join(model_dir, f"{found_model_name}-v{version}") + return model_dir + + + model_subfolders = _get_model_subfolders(model_dir, _model_regex) + + if len(model_subfolders) == 0: + if must_exist: + raise ValueError(f"No model folders found for {model_search_name}\nSearched in {model_dir}") + else: + if version is None: + version = "v1" + + model_dir = os.path.join(model_dir, f"{found_model_name}-v{version}") + return model_dir + + models_and_versions = [] + for f in model_subfolders: + d = _model_regex.match(f).groupdict() + if d["version"] is not None: + models_and_versions.append((f, normalize_version(d["version"]))) + + matching_models = [(m, v) for m, v in models_and_versions if v == version] + + if len(matching_models) == 1: + model_dir = os.path.join(model_dir, matching_models[0][0]) + return model_dir + elif len(matching_models) > 1: + raise ValueError(f"Multiple directories found for {model_search_name} with version {version}\nSearched in {model_dir}") + + top_level_has_specific_file = specific_file is not None and os.path.isfile(os.path.join(model_dir, specific_file)) + + if top_level_has_specific_file: + return model_dir + + if len(matching_models) == 0: + if must_exist and version is not None: + raise ValueError(f"No model folders found for {model_search_name} with version {version}\nSearched in {model_dir}") + + found_model_subfolder, subfolder_version = max(models_and_versions, key=lambda x: tokenize_version(x[1])) + + return os.path.join(model_dir, found_model_subfolder) + + def get_detector_model_file_path(model_name, must_exist=True): return _get_model_path(model_name, prefix="Detectors/densities", suffix=".dat", is_file=True, must_exist=must_exist) @@ -604,7 +708,31 @@ def load_flux(model_name, *args, **kwargs): def load_detector(model_name, *args, **kwargs): - return load_resource("detector", model_name, *args, **kwargs) + resource_type = "detector" + resource_name = model_name + folder = _resource_folder_by_name[resource_type] + specific_file = f"{resource_type}.py" + + abs_dir = _get_model_path(resource_name, prefix=folder, is_file=False, must_exist=True, specific_file=specific_file) + + script_fname = os.path.join(abs_dir, f"{resource_type}.py") + if os.path.isfile(script_fname): + resource_module = load_module(f"siren-{resource_type}-{resource_name}", script_fname, persist=False) + loader = getattr(resource_module, f"load_{resource_type}") + resource = loader(*args, **kwargs) + return resource + + densities_fname = os.path.join(abs_dir, "densities.dat") + materials_fname = os.path.join(abs_dir, "materials.dat") + + if os.path.isfile(densities_fname) and os.path.isfile(materials_fname): + from . import detector as _detector + detector_model = _detector.DetectorModel() + detector_model.LoadMaterialModel(materials_fname) + detector_model.LoadDetectorModel(densities_fname) + return detector_model + + raise ValueError("Could not find detector loading script \"{script_fname}\" or densities and materials files \"{densities_fname}\", \"materials_fname\"") def load_processes(model_name, *args, **kwargs):