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
"Initalize values" max_epochs = 200 batch_size = 128 lr = 0.0001 #Instead of passing in spikes, pass in smoothed spikes to do PCA on A, C, R, X, Z, Path, possible_paths, leaf_path, leaf_nodes = init.initialize(Ysmooth, D_in, K, max_epochs, batch_size, lr) Qstart = np.repeat(np.eye(D_in)[:, :, na], K, axis=2) Sstart = np.eye(D_out) #Have to pass in bern=True for spike trains kwargs = {'D_in': D_in, 'D_out': D_out, 'K': K, 'dynamics': A, 'dynamics_noise': Qstart, 'emission': C, 'emission_noise': Sstart, 'hyper_planes': R, 'possible_paths': possible_paths, 'leaf_path': leaf_path, 'leaf_nodes': leaf_nodes, 'scale': 0.01, 'bern':True} trslds = TroSLDS(**kwargs) #Instantiiate the model #Add data to model for idx in range(len(Y)): trslds._add_data(X[idx], Y[idx], Z[idx], Path[idx]) #Visualize initalization of latent states to see if it looks okay fig = plt.figure() ax = fig.add_subplot(111) for idx in range(len(X)): ax.plot(X[idx][0, 1:], X[idx][1, 1:]) # In[]: no_samples = 500 #For ICLR we ran for 1,000 samples but it converges rather quickly. 100 should be fine.
R_par = np.zeros((D_in + 1, 1)) R_par[0, 0] = 100 R_par[1, 0] = 100 r_par = np.array([0.0]) R = [] R.append(R_par) R_temp = np.zeros((D_in + 1, 2)) R_temp[:-1, 0] = np.array([-100, 100]) # 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 # In[2]: "Check to see if plotting code works" fig = plt.figure(figsize=(10, 10)) ax = fig.add_subplot(111) xmin = -15 xmax = 15 ymin = -15 ymax = 15 delta = 0.1 transform = np.hstack((np.eye(2), np.zeros((2, 1)))) X, Y, arrows = plotting.rot_vector_field(true_model.Aleaf, true_model.R, xmin, xmax, ymin, ymax, delta, depth, leaf_path, K, transform) norm = np.sqrt(arrows[:, :, 0] ** 2 + arrows[:, :, 1] ** 2)