Added verif -e to allow subsetting ensemble members

This closes #17

wxgen/database.py

@@ -34,10 +34,11 @@ class Database(object):
       _data_agg (np.array): A 3D array of data with dimensions (lead_time, variable, member*time)
       _data_matching (np.array): A 3D array of data with dimensions (lead_time, variable, member*time)
    """
-   def __init__(self, model=None):
+   def __init__(self, model=None, members=None):
       self._data_agg_cache = None
       self._data_matching_cache = None
       self.wavelet_levels = 0
+      self.members = members
       if model is None:
          self.model = wxgen.climate_model.Bin(10)
       else:
@@ -265,13 +266,13 @@ class Netcdf(Database):
    where variable_name is one or more names of weather variables. Forecast_reference_time is
    optional, i.e. both the variable and dimension could be missing.
    """
-   def __init__(self, filename, vars=None, model=None):
+   def __init__(self, filename, vars=None, model=None, members=None):
       """
       Arguments:
          filename (str): Load data from this file
          vars (list): List of indices for which variables to use
       """
-      Database.__init__(self, model)
+      Database.__init__(self, model, members)
       self.fullname = filename
       self._file = netCDF4.Dataset(filename)
 
@@ -302,9 +303,14 @@ def __init__(self, filename, vars=None, model=None):
 
       # Load data
       self.length = self._file.dimensions["time"].size
-      self._members = self._file.dimensions["ensemble_member"].size
+
+      # Determine which members to use
       V = len(self.variables)
-      M = self._members
+      if self.members is None:
+         num = self._file.dimensions["ensemble_member"].size
+         self.members = range(num)
+      M = len(self.members)
+
       has_frt = True
       if "forecast_reference_time" in self._file.dimensions:
          D = self._file.dimensions["forecast_reference_time"].size
@@ -363,18 +369,19 @@ def __init__(self, filename, vars=None, model=None):
          if var.name == "precipitation_amount":
             temp[temp < 0] = 0
          index = 0
-         for d in range(0, D):
+         for d in range(D):
             for m in range(0, M):
+               Im = self.members[m]
                if is_spatial:
                   if has_frt:
-                     self._data[:, :, :, v, index] = temp[Itimes[d], :, m, :, :]
+                     self._data[:, :, :, v, index] = temp[Itimes[d], :, Im, :, :]
                   else:
-                     self._data[:, :, :, v, index] = temp[:, m, :, :]
+                     self._data[:, :, :, v, index] = temp[:, Im, :, :]
                else:
                   if has_frt:
-                     self._data[:, :, :, v, index] = np.reshape(temp[Itimes[d], :, m], [T, Y, X])
+                     self._data[:, :, :, v, index] = np.reshape(temp[Itimes[d], :, Im], [T, Y, X])
                   else:
-                     self._data[:, :, :, v, index] = np.reshape(temp[:, m], [T, Y, X])
+                     self._data[:, :, :, v, index] = np.reshape(temp[:, Im], [T, Y, X])
                index = index + 1
 
       # If one or more values are missing for a member, set all values to nan
@@ -384,7 +391,7 @@ def __init__(self, filename, vars=None, model=None):
             self._data[:, :, :, :, e] = np.nan
             wxgen.util.debug("Removing member %d because of %d missing values" % (e, NM))
 
-      self.inittimes = np.repeat(times, self._members)
+      self.inittimes = np.repeat(times, M)
 
       self.climate_states = self.model.get(self.inittimes)
       self._file.close()

wxgen/driver.py

@@ -126,9 +126,9 @@ def main(argv):
       truth = None
       sims = None
       if args.truth is not None:
-         truth = wxgen.database.Netcdf(args.truth, None)
+         truth = wxgen.database.Netcdf(args.truth, None, members=args.members)
       if args.files is not None:
-         sims = [wxgen.database.Netcdf(file, None) for file in args.files]
+         sims = [wxgen.database.Netcdf(file, None, members=args.members) for file in args.files]
       plot.plot(sims, truth)
 
 
@@ -194,6 +194,7 @@ def get_parsers():
    sp["verif"].add_argument('-cmap', help="Colormap (e.g. jet, RdBu, Blues_r)")
    sp["verif"].add_argument('-tm', type=int, help="Time modulus (in days)", dest="timemod")
    sp["verif"].add_argument('-ts', default=1, type=int, help="Time scale (in days)", dest="timescale")
+   sp["verif"].add_argument('-e', type=wxgen.util.parse_numbers, help="Only use these members", dest="members")
 
    """
    Common options

wxgen/tests/integration_test.py

@@ -90,6 +90,10 @@ def test_truth(self):
       self.run_with_output("wxgen truth -db examples/database.nc -n 1 -t 365")
       self.run_with_output("wxgen truth -db examples/database.nc -n 2 -t 10")
 
+   def test_member_subset(self):
+      sim_filename = self.run_with_output("wxgen sim -db examples/database.nc -n 5 -t 20")
+      self.run_with_image("wxgen verif %s -m timeseries -e 0,2:3" % sim_filename)
+
    def test_README(self):
       sim_filename = self.run_with_output("wxgen sim -db examples/database.nc -n 10 -t 100")
       file = netCDF4.Dataset(sim_filename, 'r')