Skip to content

Commit

Permalink
MarkovProcess object in distribution.py, towards #639
Browse files Browse the repository at this point in the history
  • Loading branch information
sbenthall committed Jan 5, 2021
1 parent a1e16c0 commit 423b161
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 3 deletions.
2 changes: 1 addition & 1 deletion Documentation/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ Release Data: TBD
* Adds a constructor for LogNormal distributions from mean and standard deviation (#891)[https://github.com/econ-ark/HARK/pull/891/]
* Uses new LogNormal constructor in ConsPortfolioModel (#891)[https://github.com/econ-ark/HARK/pull/891/]
* calcExpectations method for taking the expectation of a distribution over a function (#884)[https://github.com/econ-ark/HARK/pull/884/]

* MarkovProcess class (#902)[https://github.com/econ-ark/HARK/pull/902]

#### Minor Changes

Expand Down
51 changes: 51 additions & 0 deletions HARK/distribution.py
Original file line number Diff line number Diff line change
Expand Up @@ -1105,3 +1105,54 @@ def calcExpectation(dstn,func=None,values=None):
# Compute expectations over the shocks and return it
f_exp = np.dot(f_query, dstn.pmf)
return f_exp


class MarkovProcess(Distribution):
"""
A representation of a discrete Markov process.
Parameters
----------
transition_matrix : np.array
An array of floats representing a probability mass for
each state transition.
seed : int
Seed for random number generator.
"""

transition_matrix = None

def __init__(self, transition_matrix, seed=0):
"""
Initialize a discrete distribution.
"""
self.transition_matrix = transition_matrix

# Set up the RNG
super().__init__(seed)

def draw(self, state):
"""
Draw new states fromt the transition matrix.
Parameters
----------
state : int or nd.array
The state or states (1-D array) from which to draw new states.
Returns
-------
new_state : int or nd.array
New states.
"""
def sample(s):
return self.RNG.choice(
self.transition_matrix.shape[1],
p = self.transition_matrix[s,:]
)

array_sample = np.frompyfunc(sample, 1, 1)

return array_sample(state)
26 changes: 24 additions & 2 deletions HARK/tests/test_distribution.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

class DiscreteDistributionTests(unittest.TestCase):
"""
Tests for simulation.py sampling distributions
Tests for distribution.py sampling distributions
with default seed.
"""

Expand All @@ -21,7 +21,7 @@ def test_drawDiscrete(self):

class DistributionClassTests(unittest.TestCase):
"""
Tests for simulation.py sampling distributions
Tests for distribution.py sampling distributions
with default seed.
"""

Expand Down Expand Up @@ -64,3 +64,25 @@ def test_Uniform(self):

def test_Bernoulli(self):
self.assertEqual(distribution.Bernoulli().draw(1)[0], False)


class MarkovProcessTests(unittest.TestCase):
"""
Tests for MarkovProcess class.
"""

def test_draw(self):

mrkv_array = np.array(
[[.75, .25],[0.1, 0.9]]
)

mp = distribution.MarkovProcess(mrkv_array)

new_state = mp.draw(np.zeros(100).astype(int))

self.assertEqual(new_state.sum(), 20)

new_state = mp.draw(new_state)

self.assertEqual(new_state.sum(), 39)

0 comments on commit 423b161

Please sign in to comment.