Source

bubble-economy / test_traders.py

Full commit
#!/usr/bin/env python
# encoding: utf-8
"""
test_traders.py

Created by dan mackinlay on 2010-10-18.
Copyright (c) 2010 __MyCompanyName__. All rights reserved.
"""

import unittest
import traders
import numpy as np
from numpy.testing import assert_array_equal
import cPickle as pickle
import gzip

class TraderSanity(unittest.TestCase):
    def setUp(self):
        self.traderset = traders.TraderSet
        traders.TraderSet(num_agents=5,dimensions=10,norm_ord=1./10.)

class BoundsCheckers(unittest.TestCase):
    def setUp(self):
        #tset could be a generic class. it's just gotta contain .vels and 
        # .locs attributes
        tset = traders.TraderSet(num_agents=4, dimensions=3)
        self.tset = tset
        
        self.vels = np.array([[ 0.,    0.25, 0.25],
                              [ 0.25, -0.25, 0.25],
                              [-0.25,  0.25, 0.25],
                              [ 0.5,  -0.5,  0.  ]])
        self.locs = np.array([[ 0.5,   0.5,  0.5], 
                              [1.25, -0.25,  0.5],
                              [1.25, -0.25,  0.5],
                              [ 0.5,  -0.5,  0. ]])
    
    def testSlideyBounds(self):
        tset = self.tset
        tset.vels = self.vels.copy()
        tset.locs = self.locs.copy()
        #use the enforce method unbound
        traders.TraderSet.enforce_slidey_boundary_conds(tset)
        self.assertTrue(np.all(tset.locs<=1.))
        self.assertTrue(np.all(tset.locs>=0.))
        
        #no change for vectors within bounds:
        assert_array_equal(tset.locs[0], self.locs[0])
        assert_array_equal(tset.vels[0], self.vels[0])
        
        #null velocities and positions in the direction of bounds
        assert_array_equal(tset.locs[1],
          np.array([1., 0., 0.5]))
        assert_array_equal(tset.vels[1],
          np.array([0., 0., 0.25]))

        #null positions only otherwise
        assert_array_equal(tset.locs[2],
          np.array([1., 0., 0.5]))
        assert_array_equal(tset.vels[2],
          np.array([-0.25, 0.25, 0.25]))
          
    def testReflectiveBounds(self):
        tset = self.tset
        tset.vels = self.vels.copy()
        tset.locs = self.locs.copy()
        
        #use the enforce method unbound
        traders.TraderSet.enforce_reflective_boundary_conds(tset)
        self.assertTrue(np.all(tset.locs<=1.))
        self.assertTrue(np.all(tset.locs>=0.))
        
        #no change for vectors within bounds:
        assert_array_equal(tset.locs[0], self.locs[0])
        assert_array_equal(tset.vels[0], self.vels[0])
        
        #reverse velocities and clip position in the direction of bounds
        assert_array_equal(tset.locs[1],
          np.array([1., 0., 0.5]))
        assert_array_equal(tset.vels[1],
          np.array([-0.25, 0.25, 0.25]))

        #clip positions only otherwise
        assert_array_equal(tset.locs[2],
          np.array([1., 0., 0.5]))
        assert_array_equal(tset.vels[2],
          np.array([-0.25, 0.25, 0.25]))
          
if __name__ == '__main__':
    unittest.main()