Added stratified sampling for precipitation.

climax/main/downscale_train.py

@@ -106,7 +106,7 @@ if __name__ == '__main__':
             # stratify training and validation dataset by number of observed
             # wet days for precipitation
             wet_days = (Obs_ds.sel(time=CALIB_PERIOD).mean(dim=('y', 'x')) >=
-                        WET_DAY_THRESHOLD).to_array().values
+                        WET_DAY_THRESHOLD).to_array().values.squeeze()
             train, valid = train_test_split(
                 CALIB_PERIOD, stratify=wet_days, test_size=0.5)
             train, valid = sorted(train), sorted(valid)  # sort chronologically