multi-basin now works

raven/models/base.py

@@ -65,7 +65,8 @@ class Raven:
                        'evspsbl': ['pet', 'evap', 'evapotranspiration'],
                        'water_volume_transport_in_river_channel': ['qobs', 'discharge', 'streamflow']
                        }
-    _parallel_parameters = ['params', 'name', 'area', 'elevation', 'latitude', 'longitude', 'region_id', 'hrus']
+    _parallel_parameters = ['params', 'nc_index', 'name', 'area', 'elevation', 'latitude', 'longitude', 'region_id',
+                            'hrus']
 
     def __init__(self, workdir=None):
         """Initialize the RAVEN model.
@@ -108,6 +109,7 @@ def __init__(self, workdir=None):
         self.exec_path = self.workdir / 'exec'
         self.final_path = self.workdir / self.final_dir
         self._psim = 0
+        self._pdim = None  # Parallel dimension (either params or nbasins)
 
     @property
     def output_path(self):
@@ -338,7 +340,7 @@ def run(self, ts, overwrite=False, **kwds):
         if isinstance(ts, (six.string_types, Path)):
             ts = [ts, ]
 
-        # Special case for parallel parameters
+        # Case for potentially parallel parameters
         pdict = {}
         for p in self._parallel_parameters:
             a = kwds.pop(p, None)
@@ -348,8 +350,11 @@ def run(self, ts, overwrite=False, **kwds):
             else:
                 pdict[p] = np.atleast_1d(a)
 
-        # Number of parallel loops is dictated by the number of parameters
-        nloops = len(pdict['params'])
+        # Number of parallel loops is dictated by the number of parameters or nc_index.
+        nloops = max(len(pdict['params']), len(pdict['nc_index']))
+        if nloops > 1:
+            self._pdim = 'nbasins' if len(pdict['nc_index']) > 1 else 'params'
+
         for key, val in pdict.items():
             if len(val) not in [1, nloops]:
                 raise ValueError("Parameter {} has incompatible dimension: {}. "
@@ -360,7 +365,7 @@ def run(self, ts, overwrite=False, **kwds):
             if len(val) == 1:
                 pdict[key] = val.repeat(nloops, axis=0)
 
-        # Update parameter objects
+        # Update non-parallel parameter objects
         for key, val in kwds.items():
 
             if key in self._rvext:
@@ -378,6 +383,7 @@ def run(self, ts, overwrite=False, **kwds):
         if self.rvi:
             self.handle_date_defaults(ts)
 
+        # Loop over parallel parameters
         procs = []
         for self.psim in range(nloops):
             for key, val in pdict.items():
@@ -445,9 +451,8 @@ def _merge_output(self, files, name):
         if name.endswith('.nc') and not isinstance(self, raven.models.RavenMultiModel):
             ds = [xr.open_dataset(fn) for fn in files]
             try:
-                # We aggregate along the params dimensions.
-                # Hard-coded for now.
-                out = xr.concat(ds, 'params', data_vars='different')
+                # We aggregate along the pdim dimensions.
+                out = xr.concat(ds, self._pdim, data_vars='different')
                 out.to_netcdf(outfn)
                 return outfn
             except (ValueError, KeyError):

tests/common.py

@@ -165,7 +165,7 @@ def _convert_2d(fn):
                 'latitude': 54.4848,
                 'longitude': -123.3659}
 
-    ds = xr.open_dataset(fn).rename({'nstations': 'region'})
+    ds = xr.open_dataset(fn, decode_times=False).rename({'nstations': 'region'})
 
     for v in ds.data_vars:
         if v not in ['lon', 'lat']:

tests/test_emulators.py

@@ -151,11 +151,13 @@ def test_parallel_basins(self, input2d):
               latitude=54.4848,
               longitude=-123.3659,
               params=[0.529, -3.396, 407.29, 1.072, 16.9, 0.947],
-              nc_index=0,
+              nc_index=[0, 0],
+              name=['basin1', 'basin2'],
               )
 
         assert len(model.diagnostics) == 2
-        assert model.hydrograph.dims['params'] == 2
+        assert len(model.hydrograph.nbasins) == 2
+        np.testing.assert_array_equal(model.hydrograph.basin_name[:], ['basin1', 'basin2'])
 
 
 class TestGR4JCN_OST: