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",