I’m trying to use a hidden Markov model for participant reponses to a behavioural psychological task. I have a basic model working for the task (code below), but am struggling to figure out how to integrate one more important feature. In my task, for some changes in the hidden state being inferred here, the participant is given a message telling them the state has changed, while in others, they are left to infer the change from the stimuli. The key parameter I want to estimate for my participants from this model is the transition probability that would cause the responses the participants made, so it’s important that I can treat these trials with the explicit signal differently. Currently, my transition probability matrix has the form [[1-p, p], [p, 1-p]]
, but for the steps that are signalled by a message, I would like it to be [[0, 1], [1, 0]]. Is there a way of having different probability matrices for different steps like this using pm.DiscreteMarkovChain()
?
with pm.Model() as model:
# Define single transition probability (simplified model)
p_switch = pm.Beta("p_switch", alpha=1.4, beta=4)
# Define prior for emission probabilities
p_cue_as_expected_for_condition = pm.Beta(
"p_cue_as_expected_for_condition", alpha=7, beta=2
)
# Build emission probability matrix
emission_probs = pm.math.stack(
[
[p_cue_as_expected_for_condition, 1 - p_cue_as_expected_for_condition],
[1 - p_cue_as_expected_for_condition, p_cue_as_expected_for_condition],
]
)
# Create a single 2x2 transition matrix
transition_probs = pm.math.stack(
[[1 - p_switch, p_switch], [p_switch, 1 - p_switch]]
)
# Initial state probabilities using Categorical distribution
start_state = 1 if df["cue_valid"].iloc[0] else 0
init_probs = np.zeros(2)
init_probs[start_state] = 0.98
init_probs[~start_state] = 0.02
init_dist = pm.Categorical.dist(p=init_probs)
# Now use DiscreteMarkovChain with static transition probabilities
hidden_states = pmx.DiscreteMarkovChain(
"hidden_states", P=transition_probs, init_dist=init_dist, shape=T
)
# Convert known_states_array to a NumPy array if it's not already
known_states_array = np.array(known_states_array)
# Find indices where hidden states are known
known_indices = np.where(~np.isnan(known_states_array))[0]
# Get the known states at those indices
known_states_values = known_states_array[known_indices].astype(int)
# Enforce known hidden states using pm.Potential
if len(known_indices) > 0:
# Create a boolean array where conditions are met
is_correct_state = pt.eq(hidden_states[known_indices], known_states_values)
# Convert boolean array to log-probabilities (0 for True, -inf for False)
logp = pt.sum(pt.switch(is_correct_state, -0.02, -3.91))
# These values are the log equivalents of the probabilities of 0.98 and 0.02
# respectively.
pm.Potential("hidden_states_observed", logp)
# Define observations
emissions = pm.Categorical(
"emissions", p=emission_probs[hidden_states], observed=observed_data
)