Implemented support for multiple users.

climax/main/preprocess/download_ERA5.py

@@ -5,18 +5,27 @@
 
 # builtins
 import os
+import pathlib
 from joblib import Parallel, delayed
 
 # externals
 import cdsapi
 import numpy as np
 
-# locals
-from climax.core.constants import ERA5_P_VARIABLES, ERA5_VARIABLES
-from climax.main.io import ERA5_PATH
+# ERA5 predictor variables on pressure levels
+ERA5_P_VARIABLES = ['geopotential', 'temperature', 'u_component_of_wind',
+                    'v_component_of_wind', 'specific_humidity']
+
+# path to store downloaded ERA-5 data
+ERA5_PATH = pathlib.Path('/mnt/CEPH_PROJECTS/FACT_CLIMAX/REANALYSIS/')
+
+# user: Daniel or Alice
+# @Alice: comment Daniel, activate Alice
+user = 'Daniel'
+# user = 'Alice'
 
 # pressure levels
-pressure_levels = ['850', '500']
+pressure_levels = ['1000', '850', '700', '500', '250']
 
 # time period
 years = [str(y) for y in np.arange(1981, 2011)]
@@ -37,6 +46,22 @@ CONFIG = {
     'area': area
 }
 
+# whether to overwrite existing files
+OVERWRITE = False
+
+# whether to skip variables on single levels
+SKIP_SINGLE_LEVELS = True
+
+# variables depending on user
+if user == 'Daniel':
+    ERA5_VARIABLES = ['geopotential', 'temperature', 'u_component_of_wind',
+                      'mean_sea_level_pressure', 'orography', '2m_temperature']
+
+if user == 'Alice':
+    ERA5_VARIABLES = ['v_component_of_wind', 'specific_humidity',
+                      'surface_pressure', '2m_dewpoint_temperature',
+                      'total_precipitation']
+
 
 if __name__ == '__main__':
 
@@ -60,18 +85,23 @@ if __name__ == '__main__':
             product = 'reanalysis-era5-pressure-levels'
             CONFIG['pressure_level'] = pressure_levels  # pressure levels
         else:
+            # whether to skip single-level variables
+            if SKIP_SINGLE_LEVELS:
+                continue
+
+            # "orography" variable was recently renamed to "geopotential": this
+            # results in conflicting file names
+            var = 'geopotential' if var == 'orography' else var
+
             # download configuration: ERA5 variable on single level
             product = 'reanalysis-era5-single-levels'
             if 'pressure_level' in CONFIG:
                 _ = CONFIG.pop('pressure_level') # remove pressure levels
 
-        # "orography" variable was recently renamed to "geopotential": this
-        # results in conflicting file names
-        var = 'geopotential' if var == 'orography' else var
-
         # split the download to the different years: CDS API cannot handle
         # requests over the whole time period
         Parallel(n_jobs=min(len(years), os.cpu_count()), verbose=51)(
             delayed(c.retrieve)(
                 product, {**CONFIG, **{'variable': var, 'year': year}}, file)
-            for file, year in zip(files, years) if not file.exists())
+            for file, year in zip(files, years) if (not file.exists() or
+            OVERWRITE))