Added group file module to Morris class and amended tests

SALib/sample/morris.py

@@ -4,7 +4,7 @@
 from . import common_args
 from ..sample import morris_oat, morris_groups, morris_optimal
 from ..util import read_param_file, scale_samples, read_group_file
-
+from collections import Iterable
 
 class Sample(object):
     '''
@@ -82,7 +82,7 @@ class Morris(Sample):
 
 
     def __init__(self, parameter_file, samples, num_levels, grid_jump, \
-                 group=None, optimal_trajectories=None):
+                 group_file=None, optimal_trajectories=None):
 
         self.parameter_file = parameter_file
         self.samples = samples
@@ -92,9 +92,8 @@ def __init__(self, parameter_file, samples, num_levels, grid_jump, \
         self.num_vars = pf['num_vars']
         self.bounds = pf['bounds']
         self.parameter_names = pf['names']
-        if group:
-          gf = read_group_file(group)
-          self.groups = gf['groups']
+        if group_file:
+          self.groups = self.compute_groups(group_file)
         else:
           self.groups = None
         self.optimal_trajectories = optimal_trajectories
@@ -118,6 +117,45 @@ def __init__(self, parameter_file, samples, num_levels, grid_jump, \
             self.create_sample_with_groups()
 
 
+    def flatten(self, l):
+        for el in l:
+            if isinstance(el, Iterable) and not isinstance(el, basestring):
+                for sub in self.flatten(el):
+                    yield sub
+            else:
+                yield el
+
+
+    def compute_groups(self, group_file):
+        gf = read_group_file(group_file)
+
+        data = gf['groups']
+
+        group_names = [g[0] for g in data]
+        param_names_from_gf = [g[1] for g in data]
+        param_names_to_check = self.flatten(param_names_from_gf)
+
+        actual_names = self.parameter_names
+
+        # Check parameter names in the group file match those from the parameter file
+        if not all([x in actual_names for x in param_names_to_check]):
+            print(group_names, param_names_from_gf, param_names_to_check)
+            raise ValueError("The parameter names from the group file do not match those from the parameter file")
+
+        # Compute the index of each parameter name and store in a dictionary
+        indices = dict([(x,i) for (i,x) in enumerate(actual_names)])
+
+        output = np.zeros((self.num_vars, len(group_names)))
+
+        # Get the data from the group file...
+        for row, group in enumerate(param_names_from_gf):
+            for param in group:
+                column = indices[param]
+                output[column,row] = 1
+
+        # ... and compile the numpy matrix
+        return np.matrix(output)
+
 
     def create_sample(self):
 

SALib/tests/test_morris.py

@@ -2,6 +2,7 @@
 from numpy.testing import assert_equal, assert_allclose
 from nose.tools import raises, with_setup
 from ..sample.morris import Morris
+from .test_util import setup_group_file
 import numpy as np
 
 
@@ -16,14 +17,6 @@ def setup_param_file():
          ofile.write("Test 3,0,1.0\n")
 
 
-def setup_group_file():
-    filename = "SALib/tests/test_group_file.txt"
-    with open(filename, "w") as ofile:
-         ofile.write("Test 1,1,0\n")
-         ofile.write("Test 2,0,1\n")
-         ofile.write("Test 3,0,1\n")
-
-
 def setup():
     setup_param_file()
     setup_group_file()
@@ -35,14 +28,14 @@ def test_group_file_read():
     Tests that a group file is read correctly
     '''
     parameter_file = "SALib/tests/test_param_file.txt"
-    group_file = "SALib/tests/test_group_file.txt"
+    group_file = "SALib/tests/test_group_file.csv"
 
     samples = 10
     num_levels = 4
     grid_jump = 2
 
     sample = Morris(parameter_file, samples, num_levels, grid_jump, \
-                 group=group_file, optimal_trajectories=None)
+                 group_file=group_file, optimal_trajectories=None)
 
     assert_equal(sample.parameter_names, ["Test 1", "Test 2", "Test 3"])
     assert_equal(sample.groups, np.matrix('1,0;0,1;0,1'))
@@ -65,7 +58,7 @@ def test_get_input_sample_scaled_works():
     grid_jump = 2
 
     sample = Morris(parameter_file, samples, num_levels, grid_jump, \
-                 group=None, optimal_trajectories=None)
+                 group_file=None, optimal_trajectories=None)
 
     sample.output_sample = np.arange(0,1.1,0.1).repeat(2).reshape((11,2))
     sample.bounds = [[10,20],[-10,10]]
@@ -96,7 +89,7 @@ def test_debug():
     grid_jump = 2
 
     sample = Morris(parameter_file, samples, num_levels, grid_jump, \
-                    group=None, optimal_trajectories=None)
+                    group_file=None, optimal_trajectories=None)
 
     desired = \
     '''Parameter File: SALib/tests/test_param_file.txt\n

SALib/tests/test_util.py

@@ -1,6 +1,6 @@
 from __future__ import division
 from numpy.testing import assert_equal, assert_allclose
-from nose.tools import raises, with_setup
+from nose.tools import raises, with_setup, eq_
 from ..util import read_param_file, scale_samples, read_group_file
 import os
 import numpy as np
@@ -28,15 +28,16 @@ def setup_tab_param_file_with_whitespace_in_names():
 
 
 def setup_group_file():
-    filename = "SALib/tests/test_group_file.txt"
+    filename = "SALib/tests/test_group_file.csv"
     with open(filename, "w") as ofile:
-         ofile.write("Group 1,Test 1\n")
-         ofile.write("Group 2,Test 2, Test 3\n")
+         ofile.write("'Group 1','Test 1'\n")
+         ofile.write("'Group 2','Test 2','Test 3'\n")
 
 
 def teardown():
     [os.remove("SALib/tests/%s" % f) for f in os.listdir("SALib/tests/") if f.endswith(".txt")]
-
+    [os.remove("SALib/tests/%s" % f) for f in os.listdir("SALib/tests/") if f.endswith(".csv")]
+    [os.remove("SALib/tests/%s" % f) for f in os.listdir("SALib/tests/") if f.endswith(".tab")]
 
 @with_setup(setup_function, teardown)
 def test_readfile():
@@ -85,13 +86,15 @@ def test_read_groupfile():
     '''
     Tests that a group file is read correctly
     '''
-    group_file = "SALib/tests/test_group_file.txt"
+    group_file = "SALib/tests/test_group_file.csv"
 
     gf = read_group_file(group_file)
 
-    assert_equal(gf['names'], ["Test 1", "Test 2", "Test 3"])
-    assert_equal(gf['groups'], np.matrix('1,0;0,1;0,1'))
+    desired = [['Group 1', ['Test 1']],['Group 2', ['Test 2', 'Test 3']]]
+    actual = gf['groups']
+    print actual
 
+    eq_(actual, desired)
 
 # Test scale samples
 def test_scale_samples():

SALib/util/__init__.py

@@ -56,23 +56,13 @@ def read_group_file(filename):
         a list of names of the variables
         a numpy matrix of factor group membership
     '''
-    names = []
-    groups = []
-    num_vars = 0
+    output = []
     num_groups = 0
 
     with open(filename) as csvfile:
-        dialect = csv.Sniffer().sniff(csvfile.read(1024))
+        dialect = csv.Sniffer().sniff(csvfile.read(8192))
         csvfile.seek(0)
         reader = csv.reader(csvfile, dialect)
         for row in reader:
-            num_vars += 1
-            names.append(row[0])
-            num_groups = 0
-            for column in row[1:]:
-                num_groups += 1
-                groups.append(int(column))
-        groups = np.array(groups).reshape(num_vars, num_groups)
-
-    return {'names': names, 'groups': np.asmatrix(groups), \
-            'num_vars': num_vars, 'num_groups': num_groups}
+          output.append([row[0], row[1:]])
+    return {'groups': output}