def make_nascar_model(): As = [ random_rotation(D_latent, np.pi / 24.), random_rotation(D_latent, np.pi / 48.) ] # Set the center points for each system centers = [np.array([+2.0, 0.]), np.array([-2.0, 0.])] bs = [ -(A - np.eye(D_latent)).dot(center) for A, center in zip(As, centers) ] # Add a "right" state As.append(np.eye(D_latent)) bs.append(np.array([+0.1, 0.])) # Add a "right" state As.append(np.eye(D_latent)) bs.append(np.array([-0.25, 0.])) # Construct multinomial regression to divvy up the space w1, b1 = np.array([+1.0, 0.0]), np.array([-2.0]) # x + b > 0 -> x > -b w2, b2 = np.array([-1.0, 0.0]), np.array([-2.0]) # -x + b > 0 -> x < b w3, b3 = np.array([0.0, +1.0]), np.array([0.0]) # y > 0 w4, b4 = np.array([0.0, -1.0]), np.array([0.0]) # y < 0 Rs = np.row_stack((100 * w1, 100 * w2, 10 * w3, 10 * w4)) r = np.concatenate((100 * b1, 100 * b2, 10 * b3, 10 * b4)) true_rslds = SLDS(D_obs, K, D_latent, transitions="recurrent_only", dynamics="gaussian", emissions="gaussian", single_subspace=True) true_rslds.dynamics.mu_init = np.array([0, 1]) true_rslds.dynamics.inv_sigma_init = np.log(1e-4) * np.ones(2) true_rslds.dynamics.As = np.array(As) true_rslds.dynamics.bs = np.array(bs) true_rslds.dynamics.inv_sigmas = np.log(1e-4) * np.ones((K, D_latent)) true_rslds.transitions.Rs = Rs true_rslds.transitions.r = r true_rslds.emissions.inv_etas = np.log(1e-2) * np.ones((1, D_obs)) return true_rslds
import matplotlib import matplotlib.pyplot as plt from ssm.models import SLDS, LDS from ssm.variational import SLDSMeanFieldVariationalPosterior, SLDSTriDiagVariationalPosterior from ssm.util import random_rotation, find_permutation # Set the parameters of the HMM T = 1000 # number of time bins K = 5 # number of discrete states D = 2 # number of latent dimensions N = 10 # number of observed dimensions # Make an SLDS with the true parameters true_slds = SLDS(N, K, D, emissions="gaussian") for k in range(K): true_slds.dynamics.As[k] = .95 * random_rotation(D, theta=(k+1) * np.pi/20) # Sample training and test data from the SLDS z, x, y = true_slds.sample(T) z_test, x_test, y_test = true_slds.sample(T) # Mask off some data mask = npr.rand(T, N) < 0.75 y_masked = y * mask # Fit an SLDS with mean field posterior print("Fitting SLDS with SVI using structured variational posterior") slds = SLDS(N, K, D, emissions="gaussian") slds.initialize(y_masked, masks=mask)
from ssm.util import random_rotation, find_permutation # In[2]: # Set the parameters of the HMM T = 1000 # number of time bins K = 5 # number of discrete states D = 2 # number of latent dimensions N = 100 # number of observed dimensions # In[3]: # Make an SLDS with the true parameters true_slds = SLDS(N, K, D, emissions="poisson_orthog", emission_kwargs=dict(link="softplus")) # Set rotational dynamics for k in range(K): true_slds.dynamics.As[k] = .95 * random_rotation( D, theta=(k + 1) * np.pi / 20) true_slds.dynamics.bs[k] = 3 * npr.randn(D) # Set an offset to make the counts larger # true_slds.emissions.ds += 10 # Sample data z, x, y = true_slds.sample(T)
true_rslds.emissions.inv_etas = np.log(1e-2) * np.ones((1, D_obs)) return true_rslds # Sample from the model true_rslds = make_nascar_model() z, x, y = true_rslds.sample(T=T) # In[4]: # Fit a robust rSLDS with its default initialization rslds = SLDS(D_obs, K, D_latent, transitions="recurrent_only", dynamics="diagonal_gaussian", emissions="gaussian_orthog", single_subspace=True) rslds.initialize(y) q = SLDSTriDiagVariationalPosterior(rslds, y) elbos = rslds.fit(q, y, num_iters=1000, initialize=False) # In[5]: # Get the posterior mean of the continuous states xhat = q.mean[0] # Find the permutation that matches the true and inferred states
npr.seed(0) import matplotlib import matplotlib.pyplot as plt from ssm.models import SLDS, LDS from ssm.util import random_rotation, find_permutation # Set the parameters of the HMM T = 1000 # number of time bins K = 5 # number of discrete states D = 2 # number of latent dimensions N = 10 # number of observed dimensions # Make an SLDS with the true parameters true_slds = SLDS(N, K, D, emissions="gaussian") for k in range(K): true_slds.dynamics.As[k] = .95 * random_rotation( D, theta=(k + 1) * np.pi / 20) z, x, y = true_slds.sample(T) z_test, x_test, y_test = true_slds.sample(T) # Mask off some data mask = npr.rand(T, N) < 0.75 # print("Fitting LDS with SVI") # lds = LDS(N, D, observations="gaussian") # lds_elbos, (lds_x, lds_x_var) = lds.fit(y * mask, masks=mask, num_iters=100) print("Fitting SLDS with SVI") slds = SLDS(N, K, D, emissions="gaussian")
true_rslds.transitions.Rs = Rs true_rslds.transitions.r = r true_rslds.emissions.inv_etas = np.log(1e-2) * np.ones((1, D_obs)) return true_rslds # Sample from the model true_rslds = make_nascar_model() z, x, y = true_rslds.sample(T=T) # Fit a robust rSLDS with its default initialization rslds = SLDS(D_obs, K, D_latent, transitions="recurrent_only", dynamics="gaussian", emissions="gaussian", single_subspace=True) elbos, (xhat, xvar) = rslds.fit(y, num_iters=1000) zhat = rslds.most_likely_states(xhat, y) # Plot some results plt.figure() plt.plot(elbos) plt.xlabel("Iteration") plt.ylabel("ELBO") plt.figure() ax1 = plt.subplot(121) plot_trajectory(z, x, ax=ax1)
true_rslds.transitions.Rs = Rs true_rslds.transitions.r = r true_rslds.emissions.inv_etas = np.log(1e-2) * np.ones((1, D_obs)) return true_rslds # Sample from the model true_rslds = make_nascar_model() z, x, y = true_rslds.sample(T=T) # Fit a robust rSLDS with its default initialization rslds = SLDS(D_obs, K, D_latent, transitions="recurrent_only", dynamics="gaussian", emissions="gaussian", single_subspace=True) # Initialize the model with the observed data. It is important # to call this before constructing the variational posterior since # the posterior constructor initialization looks at the rSLDS parameters. rslds.initialize(y) # Uncomment this to fit with stochastic variational inference instead q_svi = SLDSTriDiagVariationalPosterior(rslds, y) elbos = rslds.fit(q_svi, y, method="svi", num_iters=1000, initialize=False) xhat = q_svi.mean[0] zhat = rslds.most_likely_states(xhat, y)