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 test(shape,
plates,
axis=-1,
alpha_plates=None,
plate_axis=None,
mu=3):
if plate_axis is not None:
precomputes = [False, True]
else:
precomputes = [False]
for precompute in precomputes:
# Construct the model
D = shape[axis]
if alpha_plates is not None:
alpha = Gamma(2, 2, plates=alpha_plates)
alpha.initialize_from_random()
else:
alpha = 2
X = GaussianARD(mu, alpha, shape=shape, plates=plates)
# Some initial learning and rotator constructing
X.initialize_from_random()
Y = GaussianARD(X, 1)
Y.observe(np.random.randn(*(Y.get_shape(0))))
X.update()
if alpha_plates is not None:
alpha.update()
true_cost0_alpha = alpha.lower_bound_contribution()
rotX = RotateGaussianARD(X,
alpha,
axis=axis,
precompute=precompute)
else:
rotX = RotateGaussianARD(X,
axis=axis,
precompute=precompute)
true_cost0_X = X.lower_bound_contribution()
# Rotation matrices
I = np.identity(D)
R = np.random.randn(D, D)
if plate_axis is not None:
C = plates[plate_axis]
Q = np.random.randn(C, C)
Ic = np.identity(C)
else:
Q = None
Ic = None
# Compute bound terms
rotX.setup(plate_axis=plate_axis)
rot_cost0 = rotX.get_bound_terms(I, Q=Ic)
rot_cost1 = rotX.get_bound_terms(R, Q=Q)
self.assertAllClose(sum(rot_cost0.values()),
rotX.bound(I, Q=Ic)[0],
msg="Bound terms and total bound differ")
self.assertAllClose(sum(rot_cost1.values()),
rotX.bound(R, Q=Q)[0],
msg="Bound terms and total bound differ")
# Perform rotation
rotX.rotate(R, Q=Q)
# Check bound terms
true_cost1_X = X.lower_bound_contribution()
self.assertAllClose(true_cost1_X - true_cost0_X,
rot_cost1[X] - rot_cost0[X],
msg="Incorrect rotation cost for X")
if alpha_plates is not None:
true_cost1_alpha = alpha.lower_bound_contribution()
self.assertAllClose(
true_cost1_alpha - true_cost0_alpha,
rot_cost1[alpha] - rot_cost0[alpha],
msg="Incorrect rotation cost for alpha")
return
def test(shape,
plates,
axis=-1,
alpha_plates=None,
plate_axis=None,
mu=3):
if plate_axis is not None:
precomputes = [False, True]
else:
precomputes = [False]
for precompute in precomputes:
# Construct the model
D = shape[axis]
if alpha_plates is not None:
alpha = Gamma(3, 5, plates=alpha_plates)
alpha.initialize_from_random()
else:
alpha = 2
X = GaussianARD(mu, alpha, shape=shape, plates=plates)
# Some initial learning and rotator constructing
X.initialize_from_random()
Y = GaussianARD(X, 1)
Y.observe(np.random.randn(*(Y.get_shape(0))))
X.update()
if alpha_plates is not None:
alpha.update()
rotX = RotateGaussianARD(X,
alpha,
axis=axis,
precompute=precompute)
else:
rotX = RotateGaussianARD(X,
axis=axis,
precompute=precompute)
try:
mu.update()
except:
pass
# Rotation matrices
R = np.random.randn(D, D)
if plate_axis is not None:
C = plates[plate_axis]
Q = np.random.randn(C, C)
else:
Q = None
# Compute bound terms
rotX.setup(plate_axis=plate_axis)
if plate_axis is None:
def f_r(r):
(b, dr) = rotX.bound(np.reshape(r, np.shape(R)))
return (b, np.ravel(dr))
else:
def f_r(r):
(b, dr, dq) = rotX.bound(np.reshape(r, np.shape(R)),
Q=Q)
return (b, np.ravel(dr))
def f_q(q):
(b, dr, dq) = rotX.bound(R,
Q=np.reshape(q, np.shape(Q)))
return (b, np.ravel(dq))
# Check gradient with respect to R
err = optimize.check_gradient(f_r, np.ravel(R), verbose=False)
self.assertAllClose(err,
0,
atol=1e-4,
msg="Gradient incorrect for R")
# Check gradient with respect to Q
if plate_axis is not None:
err = optimize.check_gradient(f_q,
np.ravel(Q),
verbose=False)
self.assertAllClose(err,
0,
atol=1e-4,
msg="Gradient incorrect for Q")
return
def test(shape, plates,
axis=-1,
alpha_plates=None,
plate_axis=None,
mu=3):
if plate_axis is not None:
precomputes = [False, True]
else:
precomputes = [False]
for precompute in precomputes:
# Construct the model
D = shape[axis]
if alpha_plates is not None:
alpha = Gamma(2, 2,
plates=alpha_plates)
alpha.initialize_from_random()
else:
alpha = 2
X = GaussianARD(mu, alpha,
shape=shape,
plates=plates)
# Some initial learning and rotator constructing
X.initialize_from_random()
Y = GaussianARD(X, 1)
Y.observe(np.random.randn(*(Y.get_shape(0))))
X.update()
if alpha_plates is not None:
alpha.update()
true_cost0_alpha = alpha.lower_bound_contribution()
rotX = RotateGaussianARD(X, alpha,
axis=axis,
precompute=precompute)
else:
rotX = RotateGaussianARD(X,
axis=axis,
precompute=precompute)
true_cost0_X = X.lower_bound_contribution()
# Rotation matrices
I = np.identity(D)
R = np.random.randn(D, D)
if plate_axis is not None:
C = plates[plate_axis]
Q = np.random.randn(C, C)
Ic = np.identity(C)
else:
Q = None
Ic = None
# Compute bound terms
rotX.setup(plate_axis=plate_axis)
rot_cost0 = rotX.get_bound_terms(I, Q=Ic)
rot_cost1 = rotX.get_bound_terms(R, Q=Q)
self.assertAllClose(sum(rot_cost0.values()),
rotX.bound(I, Q=Ic)[0],
msg="Bound terms and total bound differ")
self.assertAllClose(sum(rot_cost1.values()),
rotX.bound(R, Q=Q)[0],
msg="Bound terms and total bound differ")
# Perform rotation
rotX.rotate(R, Q=Q)
# Check bound terms
true_cost1_X = X.lower_bound_contribution()
self.assertAllClose(true_cost1_X - true_cost0_X,
rot_cost1[X] - rot_cost0[X],
msg="Incorrect rotation cost for X")
if alpha_plates is not None:
true_cost1_alpha = alpha.lower_bound_contribution()
self.assertAllClose(true_cost1_alpha - true_cost0_alpha,
rot_cost1[alpha] - rot_cost0[alpha],
msg="Incorrect rotation cost for alpha")
return
def test(shape, plates,
axis=-1,
alpha_plates=None,
plate_axis=None,
mu=3):
if plate_axis is not None:
precomputes = [False, True]
else:
precomputes = [False]
for precompute in precomputes:
# Construct the model
D = shape[axis]
if alpha_plates is not None:
alpha = Gamma(3, 5,
plates=alpha_plates)
alpha.initialize_from_random()
else:
alpha = 2
X = GaussianARD(mu, alpha,
shape=shape,
plates=plates)
# Some initial learning and rotator constructing
X.initialize_from_random()
Y = GaussianARD(X, 1)
Y.observe(np.random.randn(*(Y.get_shape(0))))
X.update()
if alpha_plates is not None:
alpha.update()
rotX = RotateGaussianARD(X, alpha,
axis=axis,
precompute=precompute)
else:
rotX = RotateGaussianARD(X,
axis=axis,
precompute=precompute)
try:
mu.update()
except:
pass
# Rotation matrices
R = np.random.randn(D, D)
if plate_axis is not None:
C = plates[plate_axis]
Q = np.random.randn(C, C)
else:
Q = None
# Compute bound terms
rotX.setup(plate_axis=plate_axis)
if plate_axis is None:
def f_r(r):
(b, dr) = rotX.bound(np.reshape(r, np.shape(R)))
return (b, np.ravel(dr))
else:
def f_r(r):
(b, dr, dq) = rotX.bound(np.reshape(r, np.shape(R)),
Q=Q)
return (b, np.ravel(dr))
def f_q(q):
(b, dr, dq) = rotX.bound(R,
Q=np.reshape(q, np.shape(Q)))
return (b, np.ravel(dq))
# Check gradient with respect to R
err = optimize.check_gradient(f_r,
np.ravel(R),
verbose=False)[1]
self.assertAllClose(err, 0,
atol=1e-4,
msg="Gradient incorrect for R")
# Check gradient with respect to Q
if plate_axis is not None:
err = optimize.check_gradient(f_q,
np.ravel(Q),
verbose=False)[1]
self.assertAllClose(err, 0,
atol=1e-4,
msg="Gradient incorrect for Q")
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, :])
