Merge pull request #2 from IMMM-SFA/data-fetch

Add install supplement code for external data retrieval

.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: |

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])
+

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

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})'

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)

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

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',