[tests] Change example module tests

Add to and improve the quality of the tests for the example module. This
commit also fixes a small bug in the example module when dealing with
sparse masks to the example.rand function.

src/mdptoolbox/example.py

@@ -16,12 +16,12 @@ rand
 
 # Copyright (c) 2011-2014 Steven A. W. Cordwell
 # Copyright (c) 2009 INRA
-# 
+#
 # All rights reserved.
-# 
+#
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are met:
-# 
+#
 #   * Redistributions of source code must retain the above copyright notice,
 #     this list of conditions and the following disclaimer.
 #   * Redistributions in binary form must reproduce the above copyright notice,
@@ -30,7 +30,7 @@ rand
 #   * Neither the name of the <ORGANIZATION> nor the names of its contributors
 #     may be used to endorse or promote products derived from this software
 #     without specific prior written permission.
-# 
+#
 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
@@ -49,50 +49,50 @@ from scipy.sparse import coo_matrix, dok_matrix
 
 def forest(S=3, r1=4, r2=2, p=0.1, is_sparse=False):
     """Generate a MDP example based on a simple forest management scenario.
-    
+
     This function is used to generate a transition probability
     (``A`` × ``S`` × ``S``) array ``P`` and a reward (``S`` × ``A``) matrix
     ``R`` that model the following problem. A forest is managed by two actions:
     'Wait' and 'Cut'. An action is decided each year with first the objective
     to maintain an old forest for wildlife and second to make money selling cut
     wood. Each year there is a probability ``p`` that a fire burns the forest.
-    
+
     Here is how the problem is modelled.
     Let {0, 1 . . . ``S``-1 } be the states of the forest, with ``S``-1 being
     the oldest. Let 'Wait' be action 0 and 'Cut' be action 1.
     After a fire, the forest is in the youngest state, that is state 0.
     The transition matrix ``P`` of the problem can then be defined as follows::
-        
+
                    | p 1-p 0.......0  |
                    | .  0 1-p 0....0  |
         P[0,:,:] = | .  .  0  .       |
                    | .  .        .    |
                    | .  .         1-p |
                    | p  0  0....0 1-p |
-        
+
                    | 1 0..........0 |
                    | . .          . |
         P[1,:,:] = | . .          . |
                    | . .          . |
                    | . .          . |
                    | 1 0..........0 |
-    
+
     The reward matrix R is defined as follows::
-        
+
                  |  0  |
                  |  .  |
         R[:,0] = |  .  |
                  |  .  |
                  |  0  |
                  |  r1 |
-        
+
                  |  0  |
                  |  1  |
         R[:,1] = |  .  |
                  |  .  |
                  |  1  |
                  |  r2 |
-    
+
     Parameters
     ---------
     S : int, optional
@@ -110,7 +110,7 @@ def forest(S=3, r1=4, r2=2, p=0.1, is_sparse=False):
     is_sparse : bool, optional
         If True, then the probability transition matrices will be returned in
         sparse format, otherwise they will be in dense format. Default: False.
-    
+
     Returns
     -------
     out : tuple
@@ -120,7 +120,7 @@ def forest(S=3, r1=4, r2=2, p=0.1, is_sparse=False):
         of ``(S, A)``. If ``is_sparse=True`` then P is a tuple of length ``A``
         where each ``P[a]`` is a scipy sparse CSR format matrix of shape
         ``(S, S)``; R remains the same as in the case of ``is_sparse=False``.
-    
+
     Examples
     --------
     >>> import mdptoolbox.example
@@ -154,12 +154,12 @@ def forest(S=3, r1=4, r2=2, p=0.1, is_sparse=False):
     True
     >>> (Rsp == R).all()
     True
-    
+
     """
     assert S > 1, "The number of states S must be greater than 1."
     assert (r1 > 0) and (r2 > 0), "The rewards must be non-negative."
     assert 0 <= p <= 1, "The probability p must be in [0; 1]."
-    # Definition of Transition matrix 
+    # Definition of Transition matrix
     if is_sparse:
         P = []
         rows = list(range(S)) * 2
@@ -188,7 +188,7 @@ def forest(S=3, r1=4, r2=2, p=0.1, is_sparse=False):
 
 def rand(S, A, is_sparse=False, mask=None):
     """Generate a random Markov Decision Process.
-    
+
     Parameters
     ----------
     S : int
@@ -201,7 +201,7 @@ def rand(S, A, is_sparse=False, mask=None):
     mask : array, optional
         Array with 0 and 1 (0 indicates a place for a zero probability), shape
         can be ``(S, S)`` or ``(A, S, S)``. Default: random.
-    
+
     Returns
     -------
     out : tuple
@@ -212,7 +212,7 @@ def rand(S, A, is_sparse=False, mask=None):
         ``A``, where each ``P[a]`` is a scipy sparse CSR format matrix of shape
         ``(S, S)`` and each ``R[a]`` is a scipy sparse csr format matrix of
         shape ``(S, 1)``.
-    
+
     Examples
     --------
     >>> import numpy, mdptoolbox.example
@@ -261,7 +261,7 @@ def rand(S, A, is_sparse=False, mask=None):
     >>> # The number of non-zero elements (nnz) in P and R are equal
     >>> Psp[1].nnz == Rsp[1].nnz
     True
-    
+
     """
     # making sure the states and actions are more than one
     assert S > 1, "The number of states S must be greater than 1."
@@ -299,7 +299,12 @@ def rand(S, A, is_sparse=False, mask=None):
                 if n == 0:
                     m[randint(0, S)] = 1
                     n = 1
-                cols = where(m)[0] # m[s, :]
+                # find the columns of the vector that have non-zero elements
+                nz = m.nonzero()
+                if len(nz) == 1:
+                    cols = nz[0]
+                else:
+                    cols = nz[1]
                 vals = random(n)
                 vals = vals / vals.sum()
                 reward = 2*random(n) - ones(n)
@@ -330,7 +335,6 @@ def rand(S, A, is_sparse=False, mask=None):
                 # Make sure that there is atleast one transition in each state
                 if m.sum() == 0:
                     m[randint(0, S)] = 1
-                    n = 1
                 P[a][s] = m * random(S)
                 P[a][s] = P[a][s] / P[a][s].sum()
                 R[a][s] = (m * (2*random(S) - ones(S, dtype=int)))

src/tests/test_examples.py

@@ -5,50 +5,183 @@ Created on Sat Aug 24 14:55:05 2013
 @author: steve
 """
 
+from nose.tools import assert_equal, assert_is_none, assert_true, \
+    assert_raises
 import numpy as np
+import scipy.sparse as sp
 
 import mdptoolbox.example
 
-from .utils import ACTIONS, STATES, P_forest, R_forest, P_rand, R_rand
-from .utils import P_rand_sparse, R_rand_sparse
+## example.forest
 
-def test_example_forest_P_shape():
-    assert (P_forest == np.array([[[0.1, 0.9, 0.0],
-                         [0.1, 0.0, 0.9],
-                         [0.1, 0.0, 0.9]],
-                        [[1, 0, 0],
-                         [1, 0, 0],
-                         [1, 0, 0]]])).all()
+def assert_spacing_equal(A, B):
+    return(assert_true((np.abs(A - B) <= np.spacing(1)).all()))
 
-def test_example_forest_R_shape():
-    assert (R_forest == np.array([[0, 0],
-                        [0, 1],
-                        [4, 2]])).all()
+class TestExampleForest(object):
+    P = np.array(
+        [[[0.1, 0.9, 0.0],
+          [0.1, 0.0, 0.9],
+          [0.1, 0.0, 0.9]],
+         [[1, 0, 0],
+          [1, 0, 0],
+          [1, 0, 0]]])
+    R = np.array(
+        [[0, 0],
+         [0, 1],
+         [4, 2]])
 
-def test_example_forest_check():
+    def test_dense_PR(self):
+        P, R = mdptoolbox.example.forest()
+        assert_equal(P.shape, self.P.shape)
+        assert_equal(R.shape, self.R.shape)
+        assert_spacing_equal(P, self.P)
+        assert_spacing_equal(R, self.R)
+
+    def test_sparse_PR(self):
+        P, R = mdptoolbox.example.forest(is_sparse=True)
+        assert_equal(len(P), len(self.P))
+        for a in range(len(self.P)):
+            assert_equal(P[a].shape, self.P[a].shape)
+            assert_equal((P[a] != sp.csr_matrix(self.P[a])).nnz, 0)
+        assert_true((R == self.R).all())
+        assert_equal(R.shape, self.R.shape)
+
+def test_example_forest_dense_check():
     P, R = mdptoolbox.example.forest(10, 5, 3, 0.2)
-    assert mdptoolbox.util.check(P, R) == None
+    assert_is_none(mdptoolbox.util.check(P, R))
+
+def test_example_forest_sparse_check():
+    P, R = mdptoolbox.example.forest(S=30, is_sparse=True)
+    assert_is_none(mdptoolbox.util.check(P, R))
+
+def test_example_forest_S_raise():
+    assert_raises(AssertionError, mdptoolbox.example.forest, S=0)
+
+def test_example_forest_r1_raise():
+    assert_raises(AssertionError, mdptoolbox.example.forest, r1=0)
+
+def test_example_forest_r1_raise():
+    assert_raises(AssertionError, mdptoolbox.example.forest, r2=0)
+
+def test_example_forest_p_low_raise():
+    assert_raises(AssertionError, mdptoolbox.example.forest, p=-1)
+
+def test_example_forest_p_high_raise():
+    assert_raises(AssertionError, mdptoolbox.example.forest, p=1.1)
+
+## example.rand
+
+class TestExampleRand(object):
+    S = 3
+    A = 2
+    P = np.array(
+        [[[0.28109922699468015, 0.4285572503528079, 0.2903435226525119],
+          [0.0, 1.0, 0.0],
+          [0.0, 1.0, 0.0]],
+        [[0.4656280088928742, 0.21500769329533384, 0.31936429781179193],
+         [0.0, 0.0, 1.0],
+         [0.19806726878845474, 0.8019327312115453, 0.0]]])
+    R = np.array(
+        [[[0.7835460015641595, 0.9273255210020586, -0.2331169623484446],
+          [0.0, -0.7192984391747097, 0.0],
+          [0.0, -0.7881847856244157, -0.0]],
+         [[-0.22702203774827612, 0.8051969510588093, -0.10010002017754482],
+          [-0.0, -0.0, 0.14039354083575928],
+          [-0.06737845428738742, -0.5111488159967945, -0.0]]])
+
+    def test_dense_PR(self):
+        np.random.seed(0)
+        P, R = mdptoolbox.example.rand(self.S, self.A)
+        assert_equal(P.shape, self.P.shape)
+        assert_equal(R.shape, self.R.shape)
+        assert_spacing_equal(P, self.P)
+        assert_spacing_equal(R, self.R)
+
+    def test_sparse_PR(self):
+        P, R = mdptoolbox.example.rand(self.S, self.A, is_sparse=True)
+        for a in range(self.A):
+            assert_equal(P[a].shape, self.P[a].shape)
+            assert_equal(R[a].shape, self.R[a].shape)
+
+    def test_dense_PR_check(self):
+        np.random.seed(0)
+        P, R = mdptoolbox.example.rand(self.S, self.A)
+        assert_is_none(mdptoolbox.util.check(P, R))
+
+    def test_sparse_PR_check(self):
+        np.random.seed(0)
+        P, R = mdptoolbox.example.rand(self.S, self.A, is_sparse=True)
+        assert_is_none(mdptoolbox.util.check(P, R))
+
+    def test_S_raise(self):
+        assert_raises(AssertionError, mdptoolbox.example.rand, S=0, A=self.A)
+
+    def test_A_raise(self):
+        assert_raises(AssertionError, mdptoolbox.example.rand, S=self.S, A=0)
 
-# exampleRand
+    def test_mask_raise_1(self):
+        mask = np.random.randint(2, size=(3, 6, 9))
+        assert_raises(AssertionError, mdptoolbox.example.rand, S=self.S,
+                      A=self.A, mask=mask)
 
-def test_example_rand_dense_P_shape():
-    assert (P_rand.shape == (ACTIONS, STATES, STATES))
+    def test_mask_raise_2(self):
+        assert_raises(TypeError, mdptoolbox.example.rand, S=self.S, A=self.A,
+                      mask=True)
 
-def test_example_rand_dense_R_shape():
-    assert (R_rand.shape == (ACTIONS, STATES, STATES))
+    def test_mask_dense_1(self):
+        mask = np.array(
+            [[1, 0, 0],
+             [0, 1, 1],
+             [1, 0, 1]])
+        P, R = mdptoolbox.example.rand(S=self.S, A=self.A, mask=mask)
+        assert_is_none(mdptoolbox.util.check(P, R))
+        P, R = mdptoolbox.example.rand(S=self.S, A=self.A, mask=mask,
+                                       is_sparse=True)
+        assert_is_none(mdptoolbox.util.check(P, R))
 
-def test_example_rand_dense_check():
-    assert mdptoolbox.util.check(P_rand, R_rand) == None
+    def test_mask_dense_2(self):
+        mask = np.array(
+            [[0, 0, 0],
+             [0, 0, 0],
+             [0, 0, 0]])
+        P, R = mdptoolbox.example.rand(S=self.S, A=self.A, mask=mask)
+        assert_is_none(mdptoolbox.util.check(P, R))
+        P, R = mdptoolbox.example.rand(S=self.S, A=self.A, mask=mask,
+                                       is_sparse=True)
+        assert_is_none(mdptoolbox.util.check(P, R))
 
-def test_example_rand_sparse_P_shape():
-    assert (len(P_rand_sparse) == ACTIONS)
-    for a in range(ACTIONS):
-        assert (P_rand_sparse[a].shape == (STATES, STATES))
+    def test_mask_dense_3(self):
+        mask = np.array(
+            [[[1, 1, 0],
+              [0, 1, 1],
+              [1, 0, 1]],
+             [[0, 0, 1],
+              [0, 1, 0],
+              [1, 0, 0]]])
+        P, R = mdptoolbox.example.rand(S=self.S, A=self.A, mask=mask)
+        assert_is_none(mdptoolbox.util.check(P, R))
+        P, R = mdptoolbox.example.rand(S=self.S, A=self.A, mask=mask,
+                                       is_sparse=True)
+        assert_is_none(mdptoolbox.util.check(P, R))
 
-def test_example_rand_sparse_R_shape():
-    assert (len(R_rand_sparse) == ACTIONS)
-    for a in range(ACTIONS):
-        assert (R_rand_sparse[a].shape == (STATES, STATES))
+    def test_mask_sparse_1(self):
+        mask = sp.csr_matrix(
+            [[1, 0, 0],
+             [0, 1, 1],
+             [1, 0, 1]])
+        P, R = mdptoolbox.example.rand(S=self.S, A=self.A, mask=mask)
+        assert_is_none(mdptoolbox.util.check(P, R))
+        P, R = mdptoolbox.example.rand(S=self.S, A=self.A, mask=mask,
+                                       is_sparse=True)
+        assert_is_none(mdptoolbox.util.check(P, R))
 
-def test_example_rand_sparse_check():
-    assert mdptoolbox.util.check(P_rand_sparse, R_rand_sparse) == None
+    def test_mask_sparse_2(self):
+        mask = np.array(
+            [[0, 0, 0],
+             [0, 0, 0],
+             [0, 0, 0]])
+        P, R = mdptoolbox.example.rand(S=self.S, A=self.A, mask=mask)
+        assert_is_none(mdptoolbox.util.check(P, R))
+        P, R = mdptoolbox.example.rand(S=self.S, A=self.A, mask=mask,
+                                       is_sparse=True)
+        assert_is_none(mdptoolbox.util.check(P, R))