From b5cc17b9d394cb6705de697a4cf9befaca24d37b Mon Sep 17 00:00:00 2001 From: emilyhunt Date: Tue, 23 Sep 2025 15:28:48 +0200 Subject: [PATCH] Major expansion to data handling --- src/hr_selection_function/__init__.py | 2 +- src/hr_selection_function/config.py | 52 ++---------- src/hr_selection_function/data.py | 116 ++++++++++++++++++++++++-- 3 files changed, 117 insertions(+), 53 deletions(-) diff --git a/src/hr_selection_function/__init__.py b/src/hr_selection_function/__init__.py index 80cef9d..4144864 100644 --- a/src/hr_selection_function/__init__.py +++ b/src/hr_selection_function/__init__.py @@ -1,3 +1,3 @@ -from .config import set_data_directory +from .data import set_data_directory from .models import HR24SelectionFunction, NStarsPredictor from .gaia import estimate_gaia_data_density, estimate_m10, estimate_m_rybizki diff --git a/src/hr_selection_function/config.py b/src/hr_selection_function/config.py index 7f5b8d6..e65dc40 100644 --- a/src/hr_selection_function/config.py +++ b/src/hr_selection_function/config.py @@ -3,57 +3,17 @@ import os from pathlib import Path +from hr_selection_function.data import set_data_directory + # Internal config storage dict (user access not supported!) _CONFIG = dict() - -# Some helpers for config things -def set_data_directory(directory: Path | str): - """User-facing function for the - - Parameters - ---------- - directory : Path | str - Directory to set. Must be a pathlib.Path or a string that can be cast to a Path. - The specified directory will be checked for validity - """ - directory, first_time = _setup_data_directory(directory) - _CONFIG["data_dir"], _CONFIG["first_run"] = directory, first_time - - -def _setup_data_directory(directory: Path | str) -> tuple[Path, bool]: - """Performs setup of the data directory to be used in the package, checking user - arguments & creating it.""" - # Initial checks - if isinstance(directory, str): - try: - directory = Path(directory) - except Exception as e: - raise ValueError( - f"Unable to cast user-specified directory '{directory}' into a path. " - "Are you sure it is a valid path on your operating system?" - ) - - if not isinstance(directory, Path): - raise ValueError("Data directory path must be a pathlib.Path or string.") - - # Make the directory - first_time = False - if not directory.exists(): - print( - "This looks like your first time running the hr_selection_function package " - "/ with this data directory. Trying to create data directory at " - f"{directory}..." - ) - directory.mkdir(parents=True) - first_time = False - - return directory, first_time - - # Setup default directory directory = os.getenv("HRSF_DATA", None) if directory is None: directory = Path.home() / ".hr_selection_function" -directory, first_time = _setup_data_directory(directory) +set_data_directory(directory) + +# Some hard defaults +_CONFIG["data_url"] = "https://example.com" # Todo diff --git a/src/hr_selection_function/data.py b/src/hr_selection_function/data.py index 296f764..d7e0575 100644 --- a/src/hr_selection_function/data.py +++ b/src/hr_selection_function/data.py @@ -1,15 +1,60 @@ import requests from hr_selection_function.config import _CONFIG from tqdm import tqdm +from pathlib import Path +from functools import wraps -def _check_data_exists(): - pass +def set_data_directory(directory: Path | str): + """User-facing function for setting the data directory used by the package. + Called at package initialization automatially, but can also be called by a user + to change the directory used for data programmatically. -def _download_data(): - data_link = "https://example.com" # Todo - write_location = _CONFIG["data_dir"] / "raw_data.zip" + Parameters + ---------- + directory : Path | str + Directory to set. Must be a pathlib.Path or a string that can be cast to a Path. + The specified directory will be checked for validity + """ + _CONFIG["data_dir"], _CONFIG["first_run"] = _check_data_directory(directory) + _CONFIG["data_already_downloaded"] = _check_data_downloaded() + + +def download_data(redownload: bool = False): + """Downloads data to be used by the package.""" + if _CONFIG["data_already_downloaded"] and redownload is False: + return + + _download_file(_CONFIG["data_url"], unzip=True, write_location=_CONFIG["data_dir"]) + _CONFIG["data_already_downloaded"] = True + + +def requires_data(func): + """Wrapper for some function func that ensures that the data directory has been + created before proceeding. + """ + @wraps(func) + def inner(*args, **kwargs): + download_data(redownload=False) + func(*args, **kwargs) + return inner + + +def _download_file( + data_link: str, + unzip: bool = True, + write_location: Path | None = None, +): + """Downloads a file at a given path. Can also unzip it.""" + if write_location is None: + write_location = _CONFIG["data_dir"] + write_location = Path(write_location) + + print(f"Downloading data for hr_selection_function to {write_location}") + + if unzip: + write_location = write_location / "_temp.zip" # Fetch initial dataset response = requests.get(data_link, stream=True) @@ -26,10 +71,69 @@ def _download_data(): for chunk in response.iter_content(chunk_size=4096): file.write(chunk) + if not unzip: + return + # Unzip data print("Unzipping data...") # Todo # Delete raw .zip file print("Removing .zip file...") - # Todo + + +def _check_data_directory(directory: Path | str) -> tuple[Path, bool]: + """Performs setup of the data directory to be used in the package, checking user + arguments & creating the folder.""" + # Initial checks + if isinstance(directory, str): + try: + directory = Path(directory) + except Exception as e: + raise ValueError( + f"Unable to cast user-specified directory '{directory}' into a path. " + "Are you sure it is a valid path on your operating system?" + ) + + if not isinstance(directory, Path): + raise ValueError("Data directory path must be a pathlib.Path or string.") + + # Make the directory + first_time = False + if not directory.exists(): + print( + "This looks like your first time running the hr_selection_function package " + "/ with this data directory. Trying to create data directory at " + f"{directory}..." + ) + directory.mkdir(parents=True) + first_time = False + + return directory, first_time + + +def _check_data_downloaded(directory: Path | None = None) -> bool: + if directory is None: + directory = _CONFIG["data_dir"] + directory = Path(directory) + + # Initial directory checks (it should always exist) + if not directory.exists(): + raise ValueError( + "Specified data path does not exist. This shouldn't be able to happen, as " + "the directory should be created during package import." + ) + if not directory.is_dir(): + raise ValueError("Specified path is not a directory.") + + # Check if all files there + if not (directory / "density_hp7.parquet").exists(): + return False + if not (directory / "mcmc_samples.parquet").exists(): + return False + if not (directory / "subsample_cuts_hp7.parquet").exists(): + return False + for i in range(250): + if not (directory / f"nstars_models/{i}.ubj").exists(): + return False + return True -- 2.51.2