def model(self, value):
mixing_coeffs, means, covariances = value
components_distribution = distrax.as_distribution(
tfp.substrates.jax.distributions.MultivariateNormalFullCovariance(loc=means,
covariance_matrix=covariances,
validate_args=True))
self._model = MixtureSameFamily(mixture_distribution=distrax.Categorical(probs=mixing_coeffs),
components_distribution=components_distribution)
A = jnp.array([[0.3, 0.4, 0.3], [0.1, 0.6, 0.3], [0.2, 0.3, 0.5]])
S1 = jnp.array([[1.1, 0], [0, 0.3]])
S2 = jnp.array([[0.3, -0.5], [-0.5, 1.3]])
S3 = jnp.array([[0.8, 0.4], [0.4, 0.5]])
cov_collection = jnp.array([S1, S2, S3]) / 60
mu_collection = jnp.array([[0.3, 0.3], [0.8, 0.5], [0.3, 0.8]])
hmm = HMM(
trans_dist=distrax.Categorical(probs=A),
init_dist=distrax.Categorical(probs=initial_probs),
obs_dist=distrax.as_distribution(
tfp.substrates.jax.distributions.MultivariateNormalFullCovariance(
loc=mu_collection, covariance_matrix=cov_collection)))
n_samples, seed = 50, 100
samples_state, samples_obs = hmm_sample(hmm, n_samples, PRNGKey(seed))
xmin, xmax = 0, 1
ymin, ymax = 0, 1.2
colors = ["tab:green", "tab:blue", "tab:red"]
fig, ax = plt.subplots()
_, color_sample = plot_2dhmm(hmm, samples_obs, samples_state, colors, ax,
xmin, xmax, ymin, ymax)
pml.savefig("hmm_lillypad_2d.pdf")
fig, ax = plt.subplots()
ax.step(range(n_samples),
def MultivariateNormal(loc, covariance_matrix):
return distrax.as_distribution(
tfd.MultivariateNormalFullCovariance(loc, covariance_matrix))