diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index cc4d4b6..ce479a5 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -27,6 +27,7 @@ jobs: run: | python -m pip install --upgrade pip pip install -r requirements.txt + python setup.py install - name: Test and generate coverage report run: | diff --git a/download.py b/download.py new file mode 100644 index 0000000..ef915c8 --- /dev/null +++ b/download.py @@ -0,0 +1,53 @@ +import enum +import os + +from benedict import benedict + +from mosartwmpy.utilities.download_data import download_data + +available_data = benedict.from_yaml('./mosartwmpy/data_manifest.yaml') + +data_list = [] +data = [] + +for i, name in enumerate(available_data.keys()): + data_list.append(name) + data.append(f""" + {i + 1}) {name} - {available_data.get(f'{name}.description')}""") + +# clear the terminal +print(chr(27) + "[2J") + +print(f""" + 🎶 Welcome to the mosartwmpy download utility! 🎵 + + Please select the data you wish to download by typing the number: +""") + +for d in data: + print(f""" + {d}""") + +print(f""" + + 0) exit + +""") +try: + user_input = int(input(""" + Please select a number and press enter: """)) +except: + pass + +if not user_input or user_input == 0 or user_input > len(data): + print(""" + + Exiting... + + """) + +else: + print("") + print("") + download_data(data_list[user_input - 1]) + diff --git a/mosartwmpy/data_manifest.yaml b/mosartwmpy/data_manifest.yaml new file mode 100644 index 0000000..1eade6e --- /dev/null +++ b/mosartwmpy/data_manifest.yaml @@ -0,0 +1,13 @@ +# listing of publically downloadable data related to mosartwmpy + +sample_input: + description: Sample input dataset that can be used for testing and development; covers 1980 - 1985. + url: https://zenodo.org/record/4537907/files/mosartwmpy_sample_input_data_1980_1985.zip?download=1 + destination: ./ + +validation: + description: Result datasets that can be used for validating the model; includes results with and without WM; covers 1981-1982. + url: https://zenodo.org/record/4539693/files/mosartwmpy_validation.zip?download=1 + destination: ./validation + +# TODO add other weather files, demand files, etc as they become ready \ No newline at end of file diff --git a/mosartwmpy/model.py b/mosartwmpy/model.py index 228faee..b19214c 100644 --- a/mosartwmpy/model.py +++ b/mosartwmpy/model.py @@ -26,6 +26,7 @@ from mosartwmpy.reservoirs.reservoirs import reservoir_release from mosartwmpy.state.state import State from mosartwmpy.update.update import update +from mosartwmpy.utilities.download_data import download_data from mosartwmpy.utilities.pretty_timer import pretty_timer from mosartwmpy.utilities.inherit_docs import inherit_docs @@ -69,8 +70,8 @@ def initialize(self, config_file_path: str) -> None: self.parameters = Parameters() # sanitize the run name self.name = sanitize_filename(self.config.get('simulation.name')).replace(" ", "_") - # setup logging and output directory - Path(f'./output/{self.name}').mkdir(parents=True, exist_ok=True) + # setup logging and output directories + Path(f'./output/{self.name}/restart_files').mkdir(parents=True, exist_ok=True) logging.basicConfig( filename=f'./output/{self.name}/mosartwmpy.log', level=self.config.get('simulation.log_level', 'INFO'), @@ -195,6 +196,10 @@ def finalize(self) -> None: # simulation is over so free memory, write data, etc return + def download_data(self, *args, **kwargs) -> None: + """Downloads data related to the model.""" + download_data(*args, **kwargs) + def get_component_name(self) -> str: # TODO include version/hash info? return f'mosartwmpy ({self.git_hash})' diff --git a/mosartwmpy/output/output.py b/mosartwmpy/output/output.py index 86dd6e3..61a788e 100644 --- a/mosartwmpy/output/output.py +++ b/mosartwmpy/output/output.py @@ -157,5 +157,5 @@ def write_restart(self): logging.info('Writing restart file.') x = self.state.to_dataframe().to_xarray() - filename = f'./output/{self.name}/{self.name}_restart_{self.current_time.year}_{self.current_time.strftime("%m")}_{self.current_time.strftime("%d")}.nc' + filename = f'./output/{self.name}/restart_files/{self.name}_restart_{self.current_time.year}_{self.current_time.strftime("%m")}_{self.current_time.strftime("%d")}.nc' x.to_netcdf(filename) diff --git a/mosartwmpy/utilities/download_data.py b/mosartwmpy/utilities/download_data.py new file mode 100644 index 0000000..bbea51d --- /dev/null +++ b/mosartwmpy/utilities/download_data.py @@ -0,0 +1,109 @@ +import os +import io +import requests +import zipfile +import logging +import sys + +from benedict import benedict + + +def download_data(dataset: str, destination: str = None, manifest: str = './mosartwmpy/data_manifest.yaml') -> None: + """Convenience wrapper for the InstallSupplement class. + + Download and unpack example data supplement from Zenodo that matches the current installed + distribution. + + Args: + dataset (str): name of the dataset to download, as found in the data_manifest.yaml + destination (str): full path to the directory in which to unpack the downloaded files; must be write enabled; defaults to the directory listed in the manifest + manifest (str): full path to the manifest yaml file describing the available downloads; defaults to the bundled data_manifest.yaml + """ + + data_dictionary = benedict(manifest, format='yaml') + + if not data_dictionary.get(dataset, None): + raise Exception(f'Dataset "{dataset}" not found in the manifest ({manifest}).') + + get = InstallSupplement(url = data_dictionary.get(f'{dataset}.url'), destination = destination if destination is not None else data_dictionary.get(f'{dataset}.destination', './')) + get.fetch_zenodo() + + +class InstallSupplement: + """Download and unpack example data supplement from Zenodo that matches the current installed + distribution. + + :param example_data_directory: Full path to the directory you wish to install + the example data to. Must be write-enabled + for the user. + + """ + + def __init__(self, url, destination): + + self.initialize_logger() + self.destination = self.valid_directory(destination) + self.url = url + + def initialize_logger(self): + """Initialize logger to stdout.""" + + # initialize logger + logger = logging.getLogger() + logger.setLevel(logging.INFO) + + # logger console handler + console_handler = logging.StreamHandler(sys.stdout) + console_handler.setLevel(logging.INFO) + console_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) + logger.addHandler(console_handler) + + @staticmethod + def close_logger(): + """Shutdown logger.""" + + # Remove logging handlers + logger = logging.getLogger() + + for handler in logger.handlers[:]: + handler.close() + logger.removeHandler(handler) + + logging.shutdown() + + def valid_directory(self, directory): + """Ensure the provided directory exists.""" + + if os.path.isdir(directory): + return directory + else: + msg = f"The write directory provided by the user does not exist: {directory}" + logging.exception(msg) + self.close_logger() + raise NotADirectoryError(msg) + + def fetch_zenodo(self): + """Download and unpack the Zenodo example data supplement for the + current distribution.""" + + # retrieve content from URL + try: + logging.info(f"Downloading example data from {self.url}") + r = requests.get(self.url) + + with zipfile.ZipFile(io.BytesIO(r.content)) as zipped: + + # extract each file in the zipped dir to the project + for f in zipped.namelist(): + logging.info("Unzipped: {}".format(os.path.join(self.destination, f))) + zipped.extract(f, self.destination) + + logging.info("Download and install complete.") + + self.close_logger() + + except requests.exceptions.MissingSchema: + msg = f"Unable to download data from {self.url}" + logging.exception(msg) + self.close_logger() + raise diff --git a/setup.py b/setup.py index 93d3486..d860ce9 100644 --- a/setup.py +++ b/setup.py @@ -15,7 +15,7 @@ def get_requirements(): setup( name='mosartwmpy', - version='0.1.0', + version='0.0.1', packages=find_packages(), url='https://github.com/IMMM-SFA/mosartwmpy', license='BSD2',