Simulate Markov Chain Transitions

Problem

Simulate a Markov chain for a given number of steps. Given a transition matrix, an initial state distribution, and the number of steps, compute the state distribution after the specified number of transitions.

Solution

import numpy as np

def simulate_markov_chain(transition_matrix: np.ndarray, initial_state: int, n_steps: int, num_simulations: int = 1000) -> np.ndarray:
    n_states = transition_matrix.shape[0]
    state_counts = np.zeros(n_states)

    for _ in range(num_simulations):
        state = initial_state
        for _ in range(n_steps):
            state = np.random.choice(n_states, p=transition_matrix[state])
        state_counts[state] += 1

    return state_counts / num_simulations

def markov_chain_distribution(transition_matrix: np.ndarray, initial_distribution: np.ndarray, n_steps: int) -> np.ndarray:
    distribution = initial_distribution.copy()
    for _ in range(n_steps):
        distribution = distribution @ transition_matrix
    return distribution

Explanation

1. Simulation approach: Start from an initial state, randomly sample the next state according to the transition probabilities, repeat for n_steps. Average over many simulations.
2. Analytical approach: Multiply the initial distribution vector by the transition matrix n times. After k steps, the distribution is pi_0 * P^k.
3. Both approaches converge to the stationary distribution for ergodic chains as n_steps grows.

Complexity

• Time: O(n_steps S) for the analytical approach; O(num_simulations n_steps * S) for simulation, where S = number of states
• Space: O(S^2) for the transition matrix, O(S) for the distribution vector