new compress function

Author: corentincarton

bin/lisflood

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 
 """
 Copyright 2019 European Union

src/lisflood/global_modules/add1.py

@@ -1043,20 +1043,14 @@ def nanCheckMap(data, filename, name):
         warnings.warn(LisfloodWarning("Warning: {} of {} land values of {} (binding: '{}') are NaN".format(is_nan.sum(), is_nan.size, filename, name)))
 
 
-def xarray_reader(path):
-    ds = xr.open_mfdataset(path+'.nc', engine='netcdf4', chunks={'time': 'auto'}, combine='by_coords')
-    variable_names = [k for k in ds.variables if len(ds.variables[k].dims) == 3]
-    if len(variable_names) > 1:
-        raise LisfloodWarning('More than one variable in dataset {}'.format(path))
-    elif len(variable_names) == 0:
-        raise LisfloodWarning('Could not find a valid variable in dataset {}'.format(path))
-    else:
-        var_name = variable_names[0]
-    da = ds[var_name]
-    return da
+def compress_xarray(data):
+    maskinfo = MaskInfo.instance()
+    mask = maskinfo.info.mask.flatten()
+    data_masked = np.compress(np.logical_not(mask), data.values)
+    return data_masked
 
 
-def extract_step_xr(da, timestep):
+def date_from_step(timestep):
 
     settings = LisSettings.instance()
     binding = settings.binding
@@ -1069,5 +1063,23 @@ def extract_step_xr(da, timestep):
     if type(currentDate) is not datetime.datetime:
         currentDate = begin + datetime.timedelta(days=(currentDate - 1) * dt_day)
 
-    data = da.sel(time=currentDate)
-    return compressArray(data, pcr=False, name=da.name)
+    return currentDate
+
+
+def xarray_reader(path):
+    ds = xr.open_mfdataset(path+'.nc', engine='netcdf4', chunks={'time': 'auto'}, combine='by_coords')
+    variable_names = [k for k in ds.variables if len(ds.variables[k].dims) == 3]
+    if len(variable_names) > 1:
+        raise LisfloodWarning('More than one variable in dataset {}'.format(path))
+    elif len(variable_names) == 0:
+        raise LisfloodWarning('Could not find a valid variable in dataset {}'.format(path))
+    else:
+        var_name = variable_names[0]
+    da = ds[var_name]
+    return da
+
+
+def extract_step_xr(da, timestep):
+    cur_date = date_from_step(timestep)
+    data = da.sel(time=cur_date)
+    return compress_xarray(data)