def test_stick_transition():
K = 3
D = 2
M = 0
model = HMM(K=K, D=D, M=M, transition='sticky', observation='gaussian', transition_kwargs=dict(alpha=1, kappa=100))
print("alpha = {}, kappa = {}".format(model.transition.alpha, model.transition.kappa))
data = torch.tensor([[2,3], [4,5], [6, 7]], dtype=torch.float64)
log_prob = model.log_likelihood(data)
print("log_prob", log_prob)
samples_z, samples_x = model.sample(10)
samples_log_prob = model.log_probability(samples_x)
print("sample log_prob", samples_log_prob)
print("model transition\n", model.transition.stationary_transition_matrix)
model1 = HMM(K=K, D=D, M=M, transition='sticky', observation='gaussian', transition_kwargs=dict(alpha=1, kappa=0.1))
print("Before fit")
print("model1 transition\n", model1.transition.stationary_transition_matrix)
losses, opt = model1.fit(samples_x)
print("After fit")
print("model1 transition\n", model1.transition.stationary_transition_matrix)
model2 = HMM(K=K, D=D, M=M, transition='sticky', observation='gaussian', transition_kwargs=dict(alpha=1, kappa=20))
print("Before fit")
print("model2 transition\n", model2.transition.stationary_transition_matrix)
losses, opt = model2.fit(samples_x)
print("After fit")
print("model2 transition\n", model2.transition.stationary_transition_matrix)
model3 = HMM(K=K, D=D, M=M, transition='sticky', observation='gaussian', transition_kwargs=dict(alpha=1, kappa=100))
print("Before fit")
print("model3 transition\n", model3.transition.stationary_transition_matrix)
losses, opt = model3.fit(samples_x)
print("After fit")
print("model3 transition\n", model3.transition.stationary_transition_matrix)
bounds = np.array([[0.0, 10.0], [0.0, 8.0], [0.0, 10.0], [0.0, 8.0]])
# model
K = 3
D = 4
M = 0
obs = TruncatedNormalObservation(K=K, D=D, M=M, bounds=bounds)
# model
model = HMM(K=K, D=D, M=M, observation=obs)
# log like
log_prob = model.log_probability(data)
print("log probability = ", log_prob)
# training
print("start training...")
num_iters = 10
losses, opt = model.fit(data, num_iters=num_iters, lr=0.001)
# sampling
samplt_T = T
print("start sampling")
sample_z, sample_x = model.sample(samplt_T)
print("start sampling based on with_noise")
sample_z2, sample_x2 = model.sample(T, with_noise=True)
