def simulate_circle(no_realizations, Tmin=400, Tmax=800): # Create the tree K = 3 depth, leaf_path, possible_paths, leaf_nodes = utils.create_balanced_binary_tree( K) # Define emission parameters C = np.zeros((D_out, D_in + 1)) C[:, :-1] = np.eye(D_out) C[0, 1] = 2 S = .001 * np.eye(D_out) # Create dynamics theta_1 = -.15 * np.pi / 2 theta_2 = -.05 * np.pi / 2 theta_3 = -.35 * np.pi / 2 A = [] A.append(np.zeros((D_in, D_in + 1, 1))) At = np.zeros((D_in, D_in + 1, 2)) At[:, :-1, 0] = np.array([[np.cos(theta_3), -np.sin(theta_3)], [np.sin(theta_3), np.cos(theta_3)]]) At[:, :, 1] *= np.nan A.append(At) At = np.zeros((D_in, D_in + 1, 4)) At[:, :, 0] *= np.nan At[:, :, 1] *= np.nan At[:, :-1, 2] = np.array([[np.cos(theta_1), -np.sin(theta_1)], [np.sin(theta_1), np.cos(theta_1)]]) At[:, :-1, 3] = np.array([[np.cos(theta_2), -np.sin(theta_2)], [np.sin(theta_2), np.cos(theta_2)]]) A.append(At) Q = np.repeat(.001 * np.eye(D_in)[:, :, na], K, axis=2) # Noise covariance # Create hyperplanes R_par = np.zeros((D_in + 1, 1)) R_par[0, 0] = 100 R_par[1, 0] = 100 R = [] R.append(R_par) R_temp = np.zeros((D_in + 1, 2)) R_temp[:, 0] *= np.nan # Left hyperplane R_temp[:-1, 1] = np.array([-100, 100]) # Right hyperplane R.append(R_temp) kwargs = { 'D_in': D_in, 'D_out': D_out, 'K': K, 'dynamics': A, 'dynamics_noise': Q, 'emission': C, 'emission_noise': S, 'hyper_planes': R, 'possible_paths': possible_paths, 'leaf_path': leaf_path, 'leaf_nodes': leaf_nodes } true_model = TroSLDS(**kwargs) # Create model # Generate data from model Xreal = [] Yreal = [] Zreal = [] starting_pts = npr.uniform(-5, 5, (D_in, no_realizations)) for reals in tqdm(range(no_realizations)): T = npr.randint(Tmin, Tmax + 1) y, x, z = true_model._generate_data(T, starting_pts[:, reals]) Xreal.append(x) Yreal.append(y) Zreal.append(z) return Xreal, Yreal, Zreal, true_model
from trslds.models import TroSLDS from numpy import newaxis as na import trslds.utils as utils import matplotlib.pyplot as plt import trslds.initialize as init import trslds.plotting as plotting import seaborn as sns color_names = ["windows blue", "leaf green", "red", "orange"] colors_leaf = sns.xkcd_palette(color_names) npr.seed(0) # In[1]: D_in = 2 D_out = 2 # Create the tree K = 4 depth, leaf_path, possible_paths, leaf_nodes = utils.create_balanced_binary_tree(K) # Define emission parameters C = np.zeros((D_out, D_in + 1)) C[:, :-1] = np.eye(D_out) C[0, 1] = 2 S = .1 * np.eye(D_out) # Create dynamics A = [] A.append(np.zeros((D_in, D_in + 1, 1))) A.append(np.zeros((D_in, D_in + 1, 2))) At = np.zeros((D_in, D_in + 1, 4)) theta_1 = -.15*np.pi/2 theta_2 = -.05*np.pi/2
vmax=1, cmap=cmap, origin='lower') ax.set_aspect('auto') ax.set_title('inferred vector field') ax.set_xlim(xlim) ax.set_ylim(ylim) ax.set_xlabel('$x_1$') ax.set_ylabel('$x_2$') # In[] #Plot vector field of root node ax = fig.add_subplot(gs[0, 2]) root_depth = 1 root_K = 1 _, root_leaf_path, _, _ = utils.create_balanced_binary_tree(root_K) X, Y, arrows = plotting.rot_vector_field(At[0], trslds.R, xmin, xmax, ymin, ymax, delta, root_depth, root_leaf_path, root_K, transform) norm = np.sqrt(arrows[:, :, 0]**2 + arrows[:, :, 1]**2) U = arrows[:, :, 0] / norm V = arrows[:, :, 1] / norm ax.streamplot(X, Y, U, V, color=np.log(norm), cmap='plasma_r') ax.set_title('vector field of root node') ax.set_xlim(xlim) ax.set_ylim(ylim) ax.set_xlabel('$x_1$') ax.set_ylabel('$x_2$') # In[]
beta = 8.0 / 3.0 def f(state, t): x, y, z = state # unpack the state vector return sigma * (y - x), x * (rho - z) - y, x * y - beta * z # derivatives t = np.arange(0.0, 50.01, .01) states = (odeint(f, 2 * pt, t) / 2).T "Plot generated trajectories from second level" Qsecond = np.zeros((3, 3, 2)) Qsecond[:, :, 0] = (Qt[:, :, 0] + Qt[:, :, 1]) / 2 Qsecond[:, :, 1] = (Qt[:, :, 2] + Qt[:, :, 3]) / 2 _, second_lp, _, _ = utils.create_balanced_binary_tree(2) xnew, znew = utils.generate_trajectory(At[1], Qsecond, trslds.R, trslds.x[2][:, 2], 2, second_lp, 2, 50000, D_in) xnew = transform[:, :-1] @ xnew + transform[:, -1][:, na] ax = fig.add_subplot(gs[0, 1], projection='3d') ax.cla() for t in range(xnew[0, :].size): ax.plot(xnew[0, t:t + 2], xnew[1, t:t + 2], xnew[2, t:t + 2], color=colors_leaf[int(znew[t])]) ax.plot(states[0, :], states[1, :], states[2, :], color="slategray",