def linear_state_space_model(D=3, N=100, M=10):
# Dynamics matrix with ARD
alpha = Gamma(1e-5, 1e-5, plates=(D, ), name='alpha')
A = Gaussian(np.zeros(D), diagonal(alpha), plates=(D, ), name='A')
# Latent states with dynamics
X = GaussianMarkovChain(
np.zeros(D), # mean of x0
1e-3 * np.identity(D), # prec of x0
A, # dynamics
np.ones(D), # innovation
n=N, # time instances
name='X')
# Mixing matrix from latent space to observation space using ARD
gamma = Gamma(1e-5, 1e-5, plates=(D, ), name='gamma')
C = Gaussian(np.zeros(D), diagonal(gamma), plates=(M, 1), name='C')
# Observation noise
tau = Gamma(1e-5, 1e-5, name='tau')
# Observations
CX = Dot(C, X.as_gaussian())
Y = Normal(CX, tau, name='Y')
return (Y, CX, X, tau, C, gamma, A, alpha)
def check(D, N, mu=None, Lambda=None, rho=None, A=None):
if mu is None:
mu = np.zeros(D)
if Lambda is None:
Lambda = np.identity(D)
if rho is None:
rho = np.ones(D)
if A is None:
A = GaussianARD(3, 5, shape=(D, ), plates=(D, ))
V = np.identity(D) + np.ones((D, D))
# Construct model
X = GaussianMarkovChain(mu,
Lambda,
A,
rho,
n=N + 1,
initialize=False)
Y = Gaussian(X, V, initialize=False)
# Posterior estimation
Y.observe(np.random.randn(*(Y.get_shape(0))))
X.update()
try:
A.update()
except:
pass
try:
mu.update()
except:
pass
try:
Lambda.update()
except:
pass
try:
rho.update()
except:
pass
# Construct rotator
rotA = RotateGaussianARD(A, axis=-1)
rotX = RotateGaussianMarkovChain(X, rotA)
rotX.setup()
# Check gradient with respect to R
R = np.random.randn(D, D)
def cost(r):
(b, dr) = rotX.bound(np.reshape(r, np.shape(R)))
return (b, np.ravel(dr))
err = optimize.check_gradient(cost, np.ravel(R), verbose=False)
self.assertAllClose(err, 0, atol=1e-5, msg="Gradient incorrect")
return
def check(D, N, K, mu=None, Lambda=None, rho=None):
if mu is None:
mu = np.zeros(D)
if Lambda is None:
Lambda = np.identity(D)
if rho is None:
rho = np.ones(D)
V = np.identity(D) + np.ones((D, D))
# Construct model
B = GaussianARD(3, 5, shape=(D, K), plates=(1, D))
S = GaussianARD(2, 4, shape=(K, ), plates=(N, 1))
A = SumMultiply('dk,k->d', B, S)
X = GaussianMarkovChain(mu,
Lambda,
A,
rho,
n=N + 1,
initialize=False)
Y = Gaussian(X, V, initialize=False)
# Posterior estimation
Y.observe(np.random.randn(N + 1, D))
X.update()
B.update()
S.update()
try:
mu.update()
except:
pass
try:
Lambda.update()
except:
pass
try:
rho.update()
except:
pass
# Construct rotator
rotB = RotateGaussianARD(B, axis=-2)
rotX = RotateVaryingMarkovChain(X, B, S, rotB)
rotX.setup()
# Check gradient with respect to R
R = np.random.randn(D, D)
def cost(r):
(b, dr) = rotX.bound(np.reshape(r, np.shape(R)))
return (b, np.ravel(dr))
err = optimize.check_gradient(cost, np.ravel(R), verbose=False)
self.assertAllClose(err, 0, atol=1e-6, msg="Gradient incorrect")
return
def run():
# Create some data
N = 500
D = 2
# Initial state
x0 = np.array([0.5, -0.5])
# Dynamics (time varying)
A0 = np.array([[.9, -.4], [.4, .9]])
A1 = np.array([[.98, -.1], [.1, .98]])
l = np.linspace(0, 1, N - 1).reshape((-1, 1, 1))
A = (1 - l) * A0 + l * A1
# Innovation covariance matrix
V = np.identity(D)
# Observation noise covariance matrix
C = np.tile(np.identity(D), (N, 1, 1))
## C0 = 10*np.array([[1, 0], [0, 1]])
## C1 = 0.01*np.array([[1, 0], [0, 1]])
## C = (1-l)**2*C0 + l**2*C1
X = np.empty((N, D))
Y = np.empty((N, D))
# Simulate data
x = x0
X[0, :] = x
Y[0, :] = x + np.random.multivariate_normal(np.zeros(D), C[0, :, :])
for n in range(N - 1):
x = np.dot(A[n, :, :], x) + np.random.multivariate_normal(
np.zeros(D), V)
X[n + 1, :] = x
Y[n + 1, :] = x + np.random.multivariate_normal(
np.zeros(D), C[n + 1, :, :])
# Invert observation noise covariance to observation precision matrices
U = np.empty((N, D, D))
UY = np.empty((N, D))
for n in range(N):
U[n, :, :] = np.linalg.inv(C[n, :, :])
UY[n, :] = np.linalg.solve(C[n, :, :], Y[n, :])
# Construct VB model
Xh = GaussianMarkovChain(np.zeros(D), np.identity(D), A, np.ones(D), n=N)
Yh = Gaussian(Xh.as_gaussian(), np.identity(D), plates=(N, ))
Yh.observe(Y)
Xh.update()
xh = Xh.u[0]
varxh = utils.diagonal(Xh.u[1]) - xh**2
#err = 2 * np.sqrt(varxh)
plt.figure(1)
plt.clf()
for d in range(D):
plt.subplot(D, 1, d)
bpplt.errorplot(xh[:, d], error=2 * np.sqrt(varxh[:, d]))
plt.plot(X[:, d], 'r-')
plt.plot(Y[:, d], '.')
def run():
# Create some data
N = 500
D = 2
# Initial state
x0 = np.array([0.5, -0.5])
# Dynamics (time varying)
A0 = np.array([[.9, -.4], [.4, .9]])
A1 = np.array([[.98, -.1], [.1, .98]])
l = np.linspace(0, 1, N-1).reshape((-1,1,1))
A = (1-l)*A0 + l*A1
# Innovation covariance matrix
V = np.identity(D)
# Observation noise covariance matrix
C = np.tile(np.identity(D), (N, 1, 1))
## C0 = 10*np.array([[1, 0], [0, 1]])
## C1 = 0.01*np.array([[1, 0], [0, 1]])
## C = (1-l)**2*C0 + l**2*C1
X = np.empty((N,D))
Y = np.empty((N,D))
# Simulate data
x = x0
X[0,:] = x
Y[0,:] = x + np.random.multivariate_normal(np.zeros(D), C[0,:,:])
for n in range(N-1):
x = np.dot(A[n,:,:],x) + np.random.multivariate_normal(np.zeros(D), V)
X[n+1,:] = x
Y[n+1,:] = x + np.random.multivariate_normal(np.zeros(D), C[n+1,:,:])
# Invert observation noise covariance to observation precision matrices
U = np.empty((N,D,D))
UY = np.empty((N,D))
for n in range(N):
U[n,:,:] = np.linalg.inv(C[n,:,:])
UY[n,:] = np.linalg.solve(C[n,:,:], Y[n,:])
# Construct VB model
Xh = GaussianMarkovChain(np.zeros(D), np.identity(D), A, np.ones(D), n=N)
Yh = Gaussian(Xh.as_gaussian(), np.identity(D), plates=(N,))
Yh.observe(Y)
Xh.update()
xh = Xh.u[0]
varxh = utils.diagonal(Xh.u[1]) - xh**2
#err = 2 * np.sqrt(varxh)
plt.figure(1)
plt.clf()
for d in range(D):
plt.subplot(D,1,d)
bpplt.errorplot(xh[:,d], error=2*np.sqrt(varxh[:,d]))
plt.plot(X[:,d], 'r-')
plt.plot(Y[:,d], '.')
def check(D, N, K, mu=None, Lambda=None, rho=None):
if mu is None:
mu = np.zeros(D)
if Lambda is None:
Lambda = np.identity(D)
if rho is None:
rho = np.ones(D)
V = np.identity(D) + np.ones((D, D))
# Construct model
B = GaussianARD(3, 5, shape=(D, K), plates=(1, D))
S = GaussianARD(2, 4, shape=(K, ), plates=(N, 1))
A = SumMultiply('dk,k->d', B, S)
X = GaussianMarkovChain(mu,
Lambda,
A,
rho,
n=N + 1,
initialize=False)
Y = Gaussian(X, V, initialize=False)
# Posterior estimation
Y.observe(np.random.randn(N + 1, D))
X.update()
B.update()
S.update()
try:
mu.update()
except:
pass
try:
Lambda.update()
except:
pass
try:
rho.update()
except:
pass
# Construct rotator
rotB = RotateGaussianARD(B, axis=-2)
rotX = RotateVaryingMarkovChain(X, B, S, rotB)
# Rotation
true_cost0 = X.lower_bound_contribution()
rotX.setup()
I = np.identity(D)
R = np.random.randn(D, D)
rot_cost0 = rotX.get_bound_terms(I)
rot_cost1 = rotX.get_bound_terms(R)
self.assertAllClose(sum(rot_cost0.values()),
rotX.bound(I)[0],
msg="Bound terms and total bound differ")
self.assertAllClose(sum(rot_cost1.values()),
rotX.bound(R)[0],
msg="Bound terms and total bound differ")
rotX.rotate(R)
true_cost1 = X.lower_bound_contribution()
self.assertAllClose(true_cost1 - true_cost0,
rot_cost1[X] - rot_cost0[X],
msg="Incorrect rotation cost for X")
return
def check(D, N, mu=None, Lambda=None, rho=None, A=None):
if mu is None:
mu = np.zeros(D)
if Lambda is None:
Lambda = np.identity(D)
if rho is None:
rho = np.ones(D)
if A is None:
A = GaussianARD(3, 5, shape=(D, ), plates=(D, ))
V = np.identity(D) + np.ones((D, D))
# Construct model
X = GaussianMarkovChain(mu,
Lambda,
A,
rho,
n=N + 1,
initialize=False)
Y = Gaussian(X, V, initialize=False)
# Posterior estimation
Y.observe(np.random.randn(*(Y.get_shape(0))))
X.update()
try:
A.update()
except:
pass
try:
mu.update()
except:
pass
try:
Lambda.update()
except:
pass
try:
rho.update()
except:
pass
# Construct rotator
rotA = RotateGaussianARD(A, axis=-1)
rotX = RotateGaussianMarkovChain(X, rotA)
# Rotation
true_cost0 = X.lower_bound_contribution()
rotX.setup()
I = np.identity(D)
R = np.random.randn(D, D)
rot_cost0 = rotX.get_bound_terms(I)
rot_cost1 = rotX.get_bound_terms(R)
self.assertAllClose(sum(rot_cost0.values()),
rotX.bound(I)[0],
msg="Bound terms and total bound differ")
self.assertAllClose(sum(rot_cost1.values()),
rotX.bound(R)[0],
msg="Bound terms and total bound differ")
rotX.rotate(R)
true_cost1 = X.lower_bound_contribution()
self.assertAllClose(true_cost1 - true_cost0,
rot_cost1[X] - rot_cost0[X],
msg="Incorrect rotation cost for X")
return
def check(D, N, K,
mu=None,
Lambda=None,
rho=None):
if mu is None:
mu = np.zeros(D)
if Lambda is None:
Lambda = np.identity(D)
if rho is None:
rho = np.ones(D)
V = np.identity(D) + np.ones((D,D))
# Construct model
B = GaussianARD(3, 5,
shape=(D,K),
plates=(1,D))
S = GaussianARD(2, 4,
shape=(K,),
plates=(N,1))
A = SumMultiply('dk,k->d', B, S)
X = GaussianMarkovChain(mu,
Lambda,
A,
rho,
n=N+1,
initialize=False)
Y = Gaussian(X,
V,
initialize=False)
# Posterior estimation
Y.observe(np.random.randn(N+1,D))
X.update()
B.update()
S.update()
try:
mu.update()
except:
pass
try:
Lambda.update()
except:
pass
try:
rho.update()
except:
pass
# Construct rotator
rotB = RotateGaussianARD(B, axis=-2)
rotX = RotateVaryingMarkovChain(X, B, S, rotB)
# Rotation
true_cost0 = X.lower_bound_contribution()
rotX.setup()
I = np.identity(D)
R = np.random.randn(D, D)
rot_cost0 = rotX.get_bound_terms(I)
rot_cost1 = rotX.get_bound_terms(R)
self.assertAllClose(sum(rot_cost0.values()),
rotX.bound(I)[0],
msg="Bound terms and total bound differ")
self.assertAllClose(sum(rot_cost1.values()),
rotX.bound(R)[0],
msg="Bound terms and total bound differ")
rotX.rotate(R)
true_cost1 = X.lower_bound_contribution()
self.assertAllClose(true_cost1 - true_cost0,
rot_cost1[X] - rot_cost0[X],
msg="Incorrect rotation cost for X")
return
def check(D, N, mu=None, Lambda=None, rho=None, A=None):
if mu is None:
mu = np.zeros(D)
if Lambda is None:
Lambda = np.identity(D)
if rho is None:
rho = np.ones(D)
if A is None:
A = GaussianARD(3, 5,
shape=(D,),
plates=(D,))
V = np.identity(D) + np.ones((D,D))
# Construct model
X = GaussianMarkovChain(mu,
Lambda,
A,
rho,
n=N+1,
initialize=False)
Y = Gaussian(X,
V,
initialize=False)
# Posterior estimation
Y.observe(np.random.randn(*(Y.get_shape(0))))
X.update()
try:
A.update()
except:
pass
try:
mu.update()
except:
pass
try:
Lambda.update()
except:
pass
try:
rho.update()
except:
pass
# Construct rotator
rotA = RotateGaussianARD(A, axis=-1)
rotX = RotateGaussianMarkovChain(X, rotA)
rotX.setup()
# Check gradient with respect to R
R = np.random.randn(D, D)
def cost(r):
(b, dr) = rotX.bound(np.reshape(r, np.shape(R)))
return (b, np.ravel(dr))
err = optimize.check_gradient(cost,
np.ravel(R),
verbose=False)[1]
self.assertAllClose(err, 0,
atol=1e-5,
msg="Gradient incorrect")
return
def check(D, N, mu=None, Lambda=None, rho=None, A=None):
if mu is None:
mu = np.zeros(D)
if Lambda is None:
Lambda = np.identity(D)
if rho is None:
rho = np.ones(D)
if A is None:
A = GaussianARD(3, 5,
shape=(D,),
plates=(D,))
V = np.identity(D) + np.ones((D,D))
# Construct model
X = GaussianMarkovChain(mu,
Lambda,
A,
rho,
n=N+1,
initialize=False)
Y = Gaussian(X,
V,
initialize=False)
# Posterior estimation
Y.observe(np.random.randn(*(Y.get_shape(0))))
X.update()
try:
A.update()
except:
pass
try:
mu.update()
except:
pass
try:
Lambda.update()
except:
pass
try:
rho.update()
except:
pass
# Construct rotator
rotA = RotateGaussianARD(A, axis=-1)
rotX = RotateGaussianMarkovChain(X, rotA)
# Rotation
true_cost0 = X.lower_bound_contribution()
rotX.setup()
I = np.identity(D)
R = np.random.randn(D, D)
rot_cost0 = rotX.get_bound_terms(I)
rot_cost1 = rotX.get_bound_terms(R)
self.assertAllClose(sum(rot_cost0.values()),
rotX.bound(I)[0],
msg="Bound terms and total bound differ")
self.assertAllClose(sum(rot_cost1.values()),
rotX.bound(R)[0],
msg="Bound terms and total bound differ")
rotX.rotate(R)
true_cost1 = X.lower_bound_contribution()
self.assertAllClose(true_cost1 - true_cost0,
rot_cost1[X] - rot_cost0[X],
msg="Incorrect rotation cost for X")
return
def check(D, N, K,
mu=None,
Lambda=None,
rho=None):
if mu is None:
mu = np.zeros(D)
if Lambda is None:
Lambda = np.identity(D)
if rho is None:
rho = np.ones(D)
V = np.identity(D) + np.ones((D,D))
# Construct model
B = GaussianARD(3, 5,
shape=(D,K),
plates=(1,D))
S = GaussianARD(2, 4,
shape=(K,),
plates=(N,1))
A = SumMultiply('dk,k->d', B, S)
X = GaussianMarkovChain(mu,
Lambda,
A,
rho,
n=N+1,
initialize=False)
Y = Gaussian(X,
V,
initialize=False)
# Posterior estimation
Y.observe(np.random.randn(N+1,D))
X.update()
B.update()
S.update()
try:
mu.update()
except:
pass
try:
Lambda.update()
except:
pass
try:
rho.update()
except:
pass
# Construct rotator
rotB = RotateGaussianARD(B, axis=-2)
rotX = RotateVaryingMarkovChain(X, B, S, rotB)
rotX.setup()
# Check gradient with respect to R
R = np.random.randn(D, D)
def cost(r):
(b, dr) = rotX.bound(np.reshape(r, np.shape(R)))
return (b, np.ravel(dr))
err = optimize.check_gradient(cost,
np.ravel(R),
verbose=False)[1]
self.assertAllClose(err, 0,
atol=1e-6,
msg="Gradient incorrect")
return
def run_dlssm(y, f, mask, D, K, maxiter):
"""
Run VB inference for linear state space model with drifting dynamics.
"""
(M, N) = np.shape(y)
# Dynamics matrix with ARD
# alpha : (D) x ()
alpha = Gamma(1e-5, 1e-5, plates=(K, ), name='alpha')
# A : (K) x (K)
A = Gaussian(
np.zeros(K),
#np.identity(K),
diagonal(alpha),
plates=(K, ),
name='A_S')
A.initialize_from_value(np.identity(K))
# rho
## rho = Gamma(1e-5,
## 1e-5,
## plates=(K,),
## name="rho")
# S : () x (N-1,K)
S = GaussianMarkovChain(np.ones(K),
1e-6 * np.identity(K),
A,
np.ones(K),
n=N - 1,
name='S')
S.initialize_from_value(1 * np.ones((N - 1, K)))
# Projection matrix of the dynamics matrix
# beta : (K) x ()
beta = Gamma(1e-5, 1e-5, plates=(K, ), name='beta')
# B : (D) x (D*K)
B = Gaussian(np.zeros(D * K),
diagonal(tile(beta, D)),
plates=(D, ),
name='B')
b = np.zeros((D, D, K))
b[np.arange(D), np.arange(D), np.zeros(D, dtype=int)] = 1
B.initialize_from_value(np.reshape(1 * b, (D, D * K)))
# A : (N-1,D) x (D)
BS = MatrixDot(B, S.as_gaussian().add_plate_axis(-1), name='BS')
# Latent states with dynamics
# X : () x (N,D)
X = GaussianMarkovChain(
np.zeros(D), # mean of x0
1e-3 * np.identity(D), # prec of x0
BS, # dynamics
np.ones(D), # innovation
n=N, # time instances
name='X',
initialize=False)
X.initialize_from_value(np.random.randn(N, D))
# Mixing matrix from latent space to observation space using ARD
# gamma : (D) x ()
gamma = Gamma(1e-5, 1e-5, plates=(D, ), name='gamma')
# C : (M,1) x (D)
C = Gaussian(np.zeros(D), diagonal(gamma), plates=(M, 1), name='C')
C.initialize_from_value(np.random.randn(M, 1, D))
# Observation noise
# tau : () x ()
tau = Gamma(1e-5, 1e-5, name='tau')
# Observations
# Y : (M,N) x ()
CX = Dot(C, X.as_gaussian())
Y = Normal(CX, tau, name='Y')
#
# RUN INFERENCE
#
# Observe data
Y.observe(y, mask=mask)
# Construct inference machine
Q = VB(Y, X, S, A, alpha, B, beta, C, gamma, tau)
#
# Run inference with rotations.
#
# Rotate the D-dimensional state space (C, X)
rotB = transformations.RotateGaussianMatrixARD(B, beta, axis='rows')
rotX = transformations.RotateDriftingMarkovChain(X, B, S, rotB)
rotC = transformations.RotateGaussianARD(C, gamma)
R_X = transformations.RotationOptimizer(rotX, rotC, D)
# Rotate the K-dimensional latent dynamics space (B, S)
rotA = transformations.RotateGaussianARD(A, alpha)
rotS = transformations.RotateGaussianMarkovChain(S, A, rotA)
rotB = transformations.RotateGaussianMatrixARD(B, beta, axis='cols')
R_S = transformations.RotationOptimizer(rotS, rotB, K)
# Iterate
for ind in range(int(maxiter / 5)):
Q.update(repeat=5)
#Q.update(X, S, A, alpha, rho, B, beta, C, gamma, tau, repeat=maxiter)
R_X.rotate()
R_S.rotate()
## R_X.rotate(
## check_bound=Q.compute_lowerbound,
## check_bound_terms=Q.compute_lowerbound_terms,
## check_gradient=True
## )
## R_S.rotate(
## check_bound=Q.compute_lowerbound,
## check_bound_terms=Q.compute_lowerbound_terms,
## check_gradient=True
## )
#
# SHOW RESULTS
#
# Mean and standard deviation of the posterior
(f_mean, f_squared) = CX.get_moments()
f_std = np.sqrt(f_squared - f_mean**2)
# Plot observations space
for m in range(M):
plt.subplot(M, 1, m + 1)
plt.plot(y[m, :], 'r.')
plt.plot(f[m, :], 'b-')
bpplt.errorplot(y=f_mean[m, :], error=2 * f_std[m, :])
def run_lssm(y, f, mask, D, maxiter):
"""
Run VB inference for linear state space model.
"""
(M, N) = np.shape(y)
#
# CONSTRUCT THE MODEL
#
# Dynamic matrix
# alpha: (D) x ()
alpha = Gamma(1e-5, 1e-5, plates=(D, ), name='alpha')
# A : (D) x (D)
A = Gaussian(np.zeros(D), diagonal(alpha), plates=(D, ), name='A')
A.initialize_from_value(np.identity(D))
# Latent states with dynamics
# X : () x (N,D)
X = GaussianMarkovChain(
np.zeros(D), # mean of x0
1e-3 * np.identity(D), # prec of x0
A, # dynamics
np.ones(D), # innovation
n=N, # time instances
name='X',
initialize=False)
X.initialize_from_value(np.random.randn(N, D))
# Mixing matrix from latent space to observation space using ARD
# gamma : (D) x ()
gamma = Gamma(1e-5, 1e-5, plates=(D, ), name='gamma')
# C : (M,1) x (D)
C = Gaussian(np.zeros(D), diagonal(gamma), plates=(M, 1), name='C')
C.initialize_from_value(np.random.randn(M, 1, D))
# Observation noise
# tau : () x ()
tau = Gamma(1e-5, 1e-5, name='tau')
# Observations
# Y : (M,N) x ()
CX = Dot(C, X.as_gaussian())
Y = Normal(CX, tau, name='Y')
#
# RUN INFERENCE
#
# Observe data
Y.observe(y, mask=mask)
# Construct inference machine
Q = VB(Y, X, A, alpha, C, gamma, tau)
# Iterate
Q.update(X, A, alpha, C, gamma, tau, repeat=maxiter)
#
# SHOW RESULTS
#
# Mean and standard deviation of the posterior
(f_mean, f_squared) = CX.get_moments()
f_std = np.sqrt(f_squared - f_mean**2)
# Plot observations space
#plt.figure()
for m in range(M):
plt.subplot(M, 1, m + 1)
plt.plot(y[m, :], 'r.')
plt.plot(f[m, :], 'b-')
bpplt.errorplot(y=f_mean[m, :], error=2 * f_std[m, :])
def run_dlssm(y, f, mask, D, K, maxiter):
"""
Run VB inference for linear state space model with drifting dynamics.
"""
(M, N) = np.shape(y)
# Dynamics matrix with ARD
# alpha : (D) x ()
alpha = Gamma(1e-5,
1e-5,
plates=(K,),
name='alpha')
# A : (K) x (K)
A = Gaussian(np.zeros(K),
#np.identity(K),
diagonal(alpha),
plates=(K,),
name='A_S')
A.initialize_from_value(np.identity(K))
# rho
## rho = Gamma(1e-5,
## 1e-5,
## plates=(K,),
## name="rho")
# S : () x (N-1,K)
S = GaussianMarkovChain(np.ones(K),
1e-6*np.identity(K),
A,
np.ones(K),
n=N-1,
name='S')
S.initialize_from_value(1*np.ones((N-1,K)))
# Projection matrix of the dynamics matrix
# beta : (K) x ()
beta = Gamma(1e-5,
1e-5,
plates=(K,),
name='beta')
# B : (D) x (D*K)
B = Gaussian(np.zeros(D*K),
diagonal(tile(beta, D)),
plates=(D,),
name='B')
b = np.zeros((D,D,K))
b[np.arange(D),np.arange(D),np.zeros(D,dtype=int)] = 1
B.initialize_from_value(np.reshape(1*b, (D,D*K)))
# A : (N-1,D) x (D)
BS = MatrixDot(B,
S.as_gaussian().add_plate_axis(-1),
name='BS')
# Latent states with dynamics
# X : () x (N,D)
X = GaussianMarkovChain(np.zeros(D), # mean of x0
1e-3*np.identity(D), # prec of x0
BS, # dynamics
np.ones(D), # innovation
n=N, # time instances
name='X',
initialize=False)
X.initialize_from_value(np.random.randn(N,D))
# Mixing matrix from latent space to observation space using ARD
# gamma : (D) x ()
gamma = Gamma(1e-5,
1e-5,
plates=(D,),
name='gamma')
# C : (M,1) x (D)
C = Gaussian(np.zeros(D),
diagonal(gamma),
plates=(M,1),
name='C')
C.initialize_from_value(np.random.randn(M,1,D))
# Observation noise
# tau : () x ()
tau = Gamma(1e-5,
1e-5,
name='tau')
# Observations
# Y : (M,N) x ()
CX = Dot(C, X.as_gaussian())
Y = Normal(CX,
tau,
name='Y')
#
# RUN INFERENCE
#
# Observe data
Y.observe(y, mask=mask)
# Construct inference machine
Q = VB(Y, X, S, A, alpha, B, beta, C, gamma, tau)
#
# Run inference with rotations.
#
# Rotate the D-dimensional state space (C, X)
rotB = transformations.RotateGaussianMatrixARD(B, beta, axis='rows')
rotX = transformations.RotateDriftingMarkovChain(X, B, S, rotB)
rotC = transformations.RotateGaussianARD(C, gamma)
R_X = transformations.RotationOptimizer(rotX, rotC, D)
# Rotate the K-dimensional latent dynamics space (B, S)
rotA = transformations.RotateGaussianARD(A, alpha)
rotS = transformations.RotateGaussianMarkovChain(S, A, rotA)
rotB = transformations.RotateGaussianMatrixARD(B, beta, axis='cols')
R_S = transformations.RotationOptimizer(rotS, rotB, K)
# Iterate
for ind in range(int(maxiter/5)):
Q.update(repeat=5)
#Q.update(X, S, A, alpha, rho, B, beta, C, gamma, tau, repeat=maxiter)
R_X.rotate()
R_S.rotate()
## R_X.rotate(
## check_bound=Q.compute_lowerbound,
## check_bound_terms=Q.compute_lowerbound_terms,
## check_gradient=True
## )
## R_S.rotate(
## check_bound=Q.compute_lowerbound,
## check_bound_terms=Q.compute_lowerbound_terms,
## check_gradient=True
## )
#
# SHOW RESULTS
#
# Mean and standard deviation of the posterior
(f_mean, f_squared) = CX.get_moments()
f_std = np.sqrt(f_squared - f_mean**2)
# Plot observations space
for m in range(M):
plt.subplot(M,1,m+1)
plt.plot(y[m,:], 'r.')
plt.plot(f[m,:], 'b-')
bpplt.errorplot(y=f_mean[m,:], error=2*f_std[m,:])
def run_lssm(y, f, mask, D, maxiter):
"""
Run VB inference for linear state space model.
"""
(M, N) = np.shape(y)
#
# CONSTRUCT THE MODEL
#
# Dynamic matrix
# alpha: (D) x ()
alpha = Gamma(1e-5,
1e-5,
plates=(D,),
name='alpha')
# A : (D) x (D)
A = Gaussian(np.zeros(D),
diagonal(alpha),
plates=(D,),
name='A')
A.initialize_from_value(np.identity(D))
# Latent states with dynamics
# X : () x (N,D)
X = GaussianMarkovChain(np.zeros(D), # mean of x0
1e-3*np.identity(D), # prec of x0
A, # dynamics
np.ones(D), # innovation
n=N, # time instances
name='X',
initialize=False)
X.initialize_from_value(np.random.randn(N,D))
# Mixing matrix from latent space to observation space using ARD
# gamma : (D) x ()
gamma = Gamma(1e-5,
1e-5,
plates=(D,),
name='gamma')
# C : (M,1) x (D)
C = Gaussian(np.zeros(D),
diagonal(gamma),
plates=(M,1),
name='C')
C.initialize_from_value(np.random.randn(M,1,D))
# Observation noise
# tau : () x ()
tau = Gamma(1e-5,
1e-5,
name='tau')
# Observations
# Y : (M,N) x ()
CX = Dot(C, X.as_gaussian())
Y = Normal(CX,
tau,
name='Y')
#
# RUN INFERENCE
#
# Observe data
Y.observe(y, mask=mask)
# Construct inference machine
Q = VB(Y, X, A, alpha, C, gamma, tau)
# Iterate
Q.update(X, A, alpha, C, gamma, tau, repeat=maxiter)
#
# SHOW RESULTS
#
# Mean and standard deviation of the posterior
(f_mean, f_squared) = CX.get_moments()
f_std = np.sqrt(f_squared - f_mean**2)
# Plot observations space
#plt.figure()
for m in range(M):
plt.subplot(M,1,m+1)
plt.plot(y[m,:], 'r.')
plt.plot(f[m,:], 'b-')
bpplt.errorplot(y=f_mean[m,:], error=2*f_std[m,:])
