Check available pressure levels.

climax/core/dataset.py

@@ -12,7 +12,7 @@ import numpy as np
 import xarray as xr
 
 # locals
-from climax.core.constants import ERA5_VARIABLES, PROJECTION, ERA5_PLEVELS
+from climax.core.constants import ERA5_VARIABLES, PROJECTION
 from pysegcnn.core.utils import search_files
 
 
@@ -43,6 +43,8 @@ class NetCDFDataset(EoDataset):
 
     def __init__(self, X, y, dim='time'):
 
+        # TODO: check if conversion to array is more efficient
+
         # NetCDF dataset containing predictor variables (ERA5)
         self.X = X
 
@@ -75,7 +77,7 @@ class ERA5Dataset(EoDataset):
         self.variables = [var for var in variables if var in ERA5_VARIABLES]
 
         # pressure levels
-        self.plevels = plevels if plevels is not None else ERA5_PLEVELS
+        self.plevels = plevels
 
     def merge(self, **kwargs):
         # search dataset for each variable in root directory
@@ -91,8 +93,12 @@ class ERA5Dataset(EoDataset):
             # check if the dataset is defined on pressure levels or single
             # levels
             if 'level' in ds.dims:
+                # check specified pressure levels to retrieve
+                levels = (self.plevels if self.plevels is not None else
+                          ds.level.values)
+
                 # iterate over pressure levels to use
-                for pl in self.plevels:
+                for pl in levels:
                     # check if pressure level is available
                     if pl not in ds.level:
                         continue