def log_partition_function(natural_params, data):
if isinstance(data, list):
return sum(map(partial(log_partition_function, natural_params), data))
log_pi, log_A, log_B = natural_params
log_alpha = log_pi
for y in data:
log_alpha = logsumexp(log_alpha[:,None] + log_A, axis=0) + log_B[:,y]
return logsumexp(log_alpha)
def entropy(self, sample=None):
"""
Compute the entropy of the variational posterior distirbution.
Recall that under the structured mean field approximation
H[q(z)q(x)] = -E_{q(z)q(x)}[log q(z) + log q(x)]
= -E_q(z)[log q(z)] - E_q(x)[log q(x)]
= H[q(z)] + H[q(x)].
That is, the entropy separates into the sum of entropies for the
discrete and continuous states.
For each one, we have
E_q(u)[log q(u)] = E_q(u) [log q(u_1) + sum_t log q(u_t | u_{t-1}) + loq q(u_t) - log Z]
= E_q(u_1)[log q(u_1)] + sum_t E_{q(u_t, u_{t-1}[log q(u_t | u_{t-1})]
+ E_q(u_t)[loq q(u_t)] - log Z
where u \in {z, x} and log Z is the log normalizer. This shows that we just need the
posterior expectations and potentials, and the log normalizer of the distribution.
Note
----
We haven't implemented the exact calculations for the continuous states yet,
so for now we're approximating the continuous state entropy via samples.
"""
# Sample the continuous states
if sample is None:
sample = self.sample_continuous_states()
else:
assert isinstance(sample, list) and len(sample) == len(self.datas)
negentropy = 0
for s, prms in zip(sample, self.params):
# 1. Compute log q(x) of samples of x
negentropy += block_tridiagonal_log_probability(
s, prms["J_diag"], prms["J_lower_diag"], prms["h"])
# 2. Compute E_{q(z)}[ log q(z) ]
log_pi0 = np.log(prms["pi0"] + 1e-16) - logsumexp(prms["pi0"])
log_Ps = np.log(prms["Ps"] + 1e-16) - logsumexp(
prms["Ps"], axis=1, keepdims=True)
(Ez, Ezzp1,
normalizer) = hmm_expected_states(prms["pi0"], prms["Ps"],
prms["log_likes"])
negentropy -= normalizer # -log Z
negentropy += np.sum(Ez[0] * log_pi0) # initial factor
negentropy += np.sum(Ez * prms["log_likes"]) # unitary factors
negentropy += np.sum(Ezzp1 * log_Ps) # pairwise factors
return -negentropy
def _discrete_entropy(self):
negentropy = 0
discrete_expectations = self.discrete_expectations
for prms, (Ez, Ezzp1, normalizer) in \
zip(self.discrete_state_params, discrete_expectations):
log_pi0 = np.log(prms["pi0"] + 1e-16) - logsumexp(prms["pi0"])
log_Ps = np.log(prms["Ps"] + 1e-16) - logsumexp(prms["Ps"], axis=1, keepdims=True)
negentropy -= normalizer # -log Z
negentropy += np.sum(Ez[0] * log_pi0) # initial factor
negentropy += np.sum(Ez * prms["log_likes"]) # unitary factors
negentropy += np.sum(Ezzp1 * log_Ps) # pairwise factors
return -negentropy
def _make_grad_hmm_normalizer(argnum, ans, pi0, Ps, ll):
# Make sure everything is C contiguous and unboxed
pi0 = to_c(pi0)
Ps = to_c(Ps)
ll = to_c(ll)
dlog_pi0 = np.zeros_like(pi0)
dlog_Ps = np.zeros_like(Ps)
dll = np.zeros_like(ll)
T, K = ll.shape
# Forward pass to get alphas
alphas = np.zeros((T, K))
forward_pass(pi0, Ps, ll, alphas)
log_Ps = np.log(Ps + LOG_EPS) - logsumexp(Ps, axis=1, keepdims=True)
grad_hmm_normalizer(log_Ps, alphas, dlog_pi0, dlog_Ps, dll)
# Compute necessary gradient
# Account for the log transformation
# df/dP = df/dlogP * dlogP/dP = df/dlogP * 1 / P
if argnum == 0:
return lambda g: g * dlog_pi0 / (pi0 + DIV_EPS)
if argnum == 1:
return lambda g: g * dlog_Ps / (Ps + DIV_EPS)
if argnum == 2:
return lambda g: g * dll
def log_transition_matrices(self, data, input, mask, tag):
T, D = data.shape
# Previous state effect
log_Ps = np.tile(self.log_Ps[None, :, :], (T - 1, 1, 1))
# Input effect
log_Ps = log_Ps + np.dot(input[1:], self.Ws.T)[:, None, :]
# Past observations effect
#Off diagonal elements of transition matrix (state switches), from past observations
log_Ps_offdiag = np.tile(
np.dot(data[:-1], self.Rs.T)[:, None, :], (1, self.K, 1))
mult_offdiag = 1 - np.tile(
np.identity(self.K)[None, :, :], (log_Ps_offdiag.shape[0], 1, 1))
#Diagonal elements of transition matrix (stickiness), from past observations
log_Ps_diag = np.tile(
np.dot(data[:-1], self.Ss.T)[:, None, :], (1, self.K, 1))
mult_diag = np.tile(
np.identity(self.K)[None, :, :], (log_Ps_diag.shape[0], 1, 1))
log_Ps = log_Ps_diag * mult_diag #Diagonal elements (stickness) from past observations
log_Ps = log_Ps + np.identity(
self.K) * self.s #Diagonal elements (stickness) bias
log_Ps = log_Ps + log_Ps_offdiag * mult_offdiag #Off diagonal elements (state switching) from past observations
log_Ps = log_Ps + (1 - np.identity(
self.K)) * self.r #Off diagonal elements (state switching) bias
return log_Ps - logsumexp(log_Ps, axis=2, keepdims=True) #Normalize
def nll_GLM_GanmorCalciumAR1(w, X, Y, hyperparams, nlfun, S=10):
"""
Negative log-likelihood for a GLM with Ganmor AR1 mixture model for calcium imaging data.
Input:
w: [D x 1] vector of GLM regression weights
X: [T x D] design matrix
Y: [T x 1] calcium fluorescence observations
hyperparams: [3 x 1] model hyperparameters: log tau, log alpha, log Gaussian variance
nlfun: [func] function handle for nonlinearity
S: [scalar] number of spikes to marginalize
return_hess: [bool] flag for returning Hessian
Output:
negative log-likelihood, gradient, and Hessian
"""
# unpack hyperparams
tau, alpha, sig2 = hyperparams
# compute AR(1) diffs
taudecay = np.exp(-1.0 / tau) # decay factor for one time bin
Y = np.pad(Y, (1, 0)) # pad Y by a time bin
Ydff = (Y[1:] - taudecay * Y[:-1]) / alpha
# compute grid of spike counts
ygrid = np.arange(0, S + 1)
# Gaussian log-likelihood terms
log_gauss_grid = -0.5 * (Ydff[:, None] - ygrid[None, :])**2 / (
sig2 / alpha**2) - 0.5 * np.log(2.0 * np.pi * sig2)
Xproj = X @ w
poissConst = gammaln(ygrid + 1)
# compute neglogli, gradient, and (optionally) Hessian
f, logf, df, ddf = nlfun(Xproj)
logPcounts = logf[:, None] * ygrid[None, :] - f[:,
None] - poissConst[None, :]
# compute log-likelihood for each time bin
logjoint = log_gauss_grid + logPcounts
logli = logsumexp(logjoint, axis=1) # log likelihood for each time bin
negL = -np.sum(logli) # negative log likelihood
# gradient
dLpoiss = (df / f)[:, None] * ygrid[
None, :] - df[:, None] # deriv of Poisson log likelihood
gwts = np.sum(np.exp(logjoint - logli[:, None]) * dLpoiss,
axis=1) # gradient weights
gradient = -X.T @ gwts
# Hessian
ddLpoiss = (ddf / f - (df / f)**2)[:, None] * ygrid[None, :] - ddf[:, None]
ddL = (ddLpoiss + dLpoiss**2)
hwts = np.sum(np.exp(logjoint - logli[:, None]) * ddL,
axis=1) - gwts**2 # hessian weights
H = -X.T @ (X * hwts[:, None])
return negL, gradient, H
def backward(self, loglikhds, cython=True):
loginit, logtrans, logobs = loglikhds
beta = []
for _logobs, _logtrans in zip(logobs, logtrans):
T = _logobs.shape[0]
_beta = np.zeros((T, self.nb_states))
if cython:
backward_cy(to_c(loginit), to_c(_logtrans), to_c(_logobs),
to_c(_beta))
else:
for k in range(self.nb_states):
_beta[T - 1, k] = 0.0
_aux = np.zeros((self.nb_states, ))
for t in range(T - 2, -1, -1):
for k in range(self.nb_states):
for j in range(self.nb_states):
_aux[j] = _logtrans[t, k, j] + _beta[
t + 1, j] + _logobs[t + 1, j]
_beta[t, k] = logsumexp(_aux)
beta.append(_beta)
return beta
def log_transition(self, x, u):
logtrans = []
for _x, _u in zip(x, u):
T = np.maximum(len(_x) - 1, 1)
_logtrans = np.tile(self.logmat[None, :, :], (T, 1, 1))
logtrans.append(_logtrans - logsumexp(_logtrans, axis=-1, keepdims=True))
return logtrans
def log_transition_matrices(self, data, input, mask, tag):
T, D = data.shape
log_Ps = np.dot(input[1:], self.Ws.T)[:, None, :] # inputs
log_Ps = log_Ps + np.dot(data[:-1], self.Rs.T)[:, None, :] # past observations
log_Ps = log_Ps + self.r # bias
log_Ps = np.tile(log_Ps, (1, self.K, 1)) # expand
return log_Ps - logsumexp(log_Ps, axis=2, keepdims=True) # normalize
def accelerate_D(doftotal, ctrl):
''' ctrl: ctrl's frequency calculation
'''
if r.topt == 1 and r.Popt == 0:
pthick = [
r.tmin / lam0 + (thick[i] - r.tmin / lam0) * doftotal[i]
for i in range(Nlayer)
]
pscale = 1.
dofold = doftotal[Nlayer:]
elif r.topt == 0 and r.Popt == 1:
pthick = thick
pscale = 1. + (r.Periodmax - r.Period) / r.Period * doftotal[0]
dofold = doftotal[1:]
elif r.topt == 1 and r.Popt == 1:
pscale = 1. + (r.Periodmax - r.Period) / r.Period * doftotal[0]
pthick = [
r.tmin / lam0 + (thick[i] - r.tmin / lam0) * doftotal[1 + i]
for i in range(Nlayer)
]
dofold = doftotal[Nlayer + 1:]
else:
pscale = 1.
pthick = thick
dofold = doftotal
# RCWA
freqcmp = freq_list[ctrl] * (1 + 1j / 2 / Qref)
planewave = {'p_amp': 0, 's_amp': 1, 'p_phase': 0, 's_phase': 0}
phi = 0.
theta = r.angle
obj, dof, epsdiff = rcwa_assembly(dofold, freqcmp, theta, phi, planewave,
pthick, pscale)
R, _ = obj.RT_Solve(normalize=1)
if r.polarization == 'ps' or r.polarization == 'sp':
planewave = {'p_amp': 1, 's_amp': 0, 'p_phase': 0, 's_phase': 0}
obj.MakeExcitationPlanewave(planewave['p_amp'],
planewave['p_phase'],
planewave['s_amp'],
planewave['s_phase'],
order=0)
R2, _ = obj.RT_Solve(normalize=1)
# the minimal of reflection
Inc = 500.
R = Inc / logsumexp(Inc / np.array([R, R2]))
# mass
rho = mload
for i in range(Nlayer):
mtmp = lam0 * pthick[i] * mstruct[i].density * np.mean(
dof[i * Nx * Ny:(i + 1) * Nx * Ny])
rho = rho + mtmp
rho = rho**r.mpower
integrand = rho / R * gamma[ctrl] * beta[ctrl] / (1 - beta[ctrl])**2
integrand = cons.c**3 / 2 / laserP * integrand * dbeta / 1e9
return integrand
def filter(self, obs, act=None):
logliklhds = self.log_likelihoods(obs, act)
alpha = self.forward(logliklhds)
belief = [
np.exp(_alpha - logsumexp(_alpha, axis=1, keepdims=True))
for _alpha in alpha
]
return belief
def mixture_log_density(var_mixture_params, x):
"""Returns a weighted average over component densities."""
log_weights, var_params = unpack_mixture_params(var_mixture_params)
component_log_densities = np.vstack(
[component_log_density(params_k, x) for params_k in var_params]).T
return logsumexp(component_log_densities + log_weights,
axis=1,
keepdims=False)
def log_transition(self, x, u):
logtrans = []
for _x, _u in zip(x, u):
T = np.maximum(len(_x) - 1, 1)
_in = np.hstack((_x[:T, :], _u[:T, :self.dm_act]))
_logtrans = to_npy(self.rnr.forward(to_torch(_in)))
logtrans.append(_logtrans - logsumexp(_logtrans, axis=-1, keepdims=True))
return logtrans
def log_transition_matrices(self, data, input, mask, tag):
T = data.shape[0]
assert input.shape[0] == T
# Previous state effect
log_Ps = np.tile(self.log_Ps[None, :, :], (T-1, 1, 1))
# Input effect
log_Ps = log_Ps + np.dot(input[1:], self.Ws.T)[:, None, :]
return log_Ps - logsumexp(log_Ps, axis=2, keepdims=True)
def transition_matrix(self, data, input, mask, tag):
# For a single data point: data is x_{t-1}, input is x_t
log_Ps = self.log_Ps[None, :, :]
# Input effect
log_Ps = log_Ps + np.dot(input, self.Ws.T)[:, None, :]
# Past observations effect
log_Ps = log_Ps + np.dot(data, self.Rs.T)[:, None, :]
return np.exp(log_Ps - logsumexp(log_Ps, axis=2, keepdims=True))
def neural_net_predict(params, inputs):
"""Implements a deep neural network for classification.
params is a list of (weights, bias) tuples.
inputs is an (N x D) matrix.
returns normalized class log-probabilities."""
for W, b in params:
outputs = np.dot(inputs, W) + b
inputs = np.tanh(outputs)
return outputs - logsumexp(outputs, axis=1, keepdims=True)
def _lowerbound(self, x, K=1):
"""
Compute the lowerbound
"""
# define network
W1, W2, W3, W4, W5 = self.params["W1"], self.params["W2"], self.params["W3"],\
self.params["W4"], self.params["W5"]
b1, b2, b3, b4, b5 = self.params["b1"], self.params["b2"], self.params["b3"],\
self.params["b4"], self.params["b5"]
if self.continuous:
W6, b6 = self.params["W6"], self.params["b6"]
activation = lambda x: np.log(1 + np.exp(x))
else:
activation = lambda x: np.tanh(x, x)
sigmoid = lambda x: 1. / (1 + np.exp(-x))
# IWAE: first replicate the input
N = x.shape[1]
#x = np.repeat(x, K, axis = 1)
#x = np.tile(x.reshape(-1, 1), (1, K)).reshape(x.shape[0], K * N)
x = np.tile(x, (1, K))
# compute forward pass
h_encoder = activation(W1.dot(x) + b1)
mu_encoder = W2.dot(h_encoder) + b2
log_sigma_encoder = 0.5 * (W3.dot(h_encoder) + b3)
eps = np.random.randn(self.n_hidden_variables, x.shape[1])
z = mu_encoder + np.exp(log_sigma_encoder) * eps
h_decoder = activation(W4.dot(z) + b4)
y = sigmoid(W5.dot(h_decoder) + b5)
if self.continuous:
log_sigma_decoder = 0.5 * (W6.dot(h_decoder) + b6)
logpxz = np.sum(-(0.5 * np.log(2 * np.pi) + log_sigma_decoder) -
0.5 * ((x - y) / np.exp(log_sigma_decoder))**2,
axis=0)
else:
logpxz = np.sum(x * np.log(y) + (1 - x) * np.log(1 - y), axis=0)
KLD = 0.5 * np.sum(1 + 2 * log_sigma_encoder - mu_encoder**2 -
np.exp(2 * log_sigma_encoder),
axis=0)
lowerbound_x = logpxz + KLD
# then compute the weights
log_ws_matrix = lowerbound_x.reshape(K, N)
# now compute the IWAE bound
lowerbound = np.sum(logsumexp(log_ws_matrix, axis=0) - np.log(K))
return lowerbound
def log_prior(self):
K = self.K
log_P = self.log_Ps - logsumexp(self.log_Ps, axis=1, keepdims=True)
lp = 0
for k in range(K):
alpha = self.alpha * np.ones(K) + self.kappa * (np.arange(K) == k)
lp += np.dot((alpha - 1), log_P[k])
return lp
def forward_pass_np(log_pi0, log_Ps, log_likes):
T, K = log_likes.shape
alphas = []
alphas.append(log_likes[0] + log_pi0)
for t in range(T-1):
anext = logsumexp(alphas[t] + log_Ps[t].T, axis=1)
anext += log_likes[t+1]
alphas.append(anext)
return np.array(alphas)
def log_transition_matrices(self, data, input, mask, tag):
T, D = data.shape
# Previous state effect
log_Ps = np.tile(self.log_Ps[None, :, :], (T-1, 1, 1))
# Input effect
log_Ps = log_Ps + np.dot(input[1:], self.Ws.T)[:, None, :]
# Past observations effect
log_Ps = log_Ps + np.dot(data[:-1], self.Rs.T)[:, None, :]
return log_Ps - logsumexp(log_Ps, axis=2, keepdims=True)
def m_step(self, expectations, datas, inputs, masks, tags, **kwargs):
K = self.K
P = sum([np.sum(Ezzp1, axis=0) for _, Ezzp1, _ in expectations]) + 1e-32
P = np.nan_to_num(P / P.sum(axis=-1, keepdims=True))
# Set rows that are all zero to uniform
P = np.where(P.sum(axis=-1, keepdims=True) == 0, 1.0 / K, P)
log_P = np.log(P)
self.log_Ps = log_P - logsumexp(log_P, axis=-1, keepdims=True)
def log_prior(self):
K = self.K
Ps = np.exp(self.log_Ps - logsumexp(self.log_Ps, axis=1, keepdims=True))
lp = 0
for k in range(K):
alpha = self.alpha * np.ones(K) + self.kappa * (np.arange(K) == k)
lp += dirichlet.logpdf(Ps[k], alpha)
return lp
def get_error_and_ll(w, v_prior, X, y, K, location, scale):
v_noise = np.exp(parser.get(w, 'log_v_noise')[0, 0]) * scale**2
q = get_parameters_q(w, v_prior)
samples_q = draw_samples(q, K)
outputs = predict(samples_q, X) * scale + location
log_factor = -0.5 * np.log(2 * math.pi * v_noise) - 0.5 * (
np.tile(y, (1, K)) - np.array(outputs))**2 / v_noise
ll = np.mean(logsumexp(log_factor - np.log(K), 1))
error = np.sqrt(np.mean((y - np.mean(outputs, 1, keepdims=True))**2))
return error, ll
def hmm_expected_states(log_pi0, log_Ps, ll, memlimit=2**31):
T, K = ll.shape
# Make sure everything is C contiguous
log_pi0 = to_c(log_pi0)
log_Ps = to_c(log_Ps)
ll = to_c(ll)
alphas = np.zeros((T, K))
forward_pass(log_pi0, log_Ps, ll, alphas)
normalizer = logsumexp(alphas[-1])
betas = np.zeros((T, K))
backward_pass(log_Ps, ll, betas)
# Compute E[z_t] for t = 1, ..., T
expected_states = alphas + betas
expected_states -= logsumexp(expected_states, axis=1, keepdims=True)
expected_states = np.exp(expected_states)
# Compute E[z_t, z_{t+1}] for t = 1, ..., T-1
# Note that this is an array of size T*K*K, which can be quite large.
# To be a bit more frugal with memory, first check if the given log_Ps
# are TxKxK. If so, instantiate the full expected joints as well, since
# we will need them for the M-step. However, if log_Ps is 1xKxK then we
# know that the transition matrix is stationary, and all we need for the
# M-step is the sum of the expected joints.
stationary = (log_Ps.shape[0] == 1)
if not stationary:
expected_joints = alphas[:-1, :, None] + betas[1:, None, :] + ll[
1:, None, :] + log_Ps
expected_joints -= expected_joints.max((1, 2))[:, None, None]
expected_joints = np.exp(expected_joints)
expected_joints /= expected_joints.sum((1, 2))[:, None, None]
else:
# Compute the sum over time axis of the expected joints
expected_joints = np.zeros((K, K))
compute_stationary_expected_joints(alphas, betas, ll, log_Ps[0],
expected_joints)
expected_joints = expected_joints[None, :, :]
return expected_states, expected_joints, normalizer
def IWELBO(params, log_p, log_q, sample_q, M_iw_train, num_copies_training):
_, lp, lq, _ = objective_utils(params, log_p, log_q, sample_q, M_iw_train,
num_copies_training)
# This will also work, but will not evaluate to IW-ELBO
# lR = lp - lq
# weights = autograd.core.getval(utils.softmax_matrix(lR))
# targets = lR
# return np.mean(np.sum(weights*targets, -1))
return np.mean(autoscipy.logsumexp(lp - lq, -1)) - np.log(M_iw_train)
def hmm_normalizer(pi0, Ps, ll):
T, K = ll.shape
alphas = np.zeros((T, K))
# Make sure everything is C contiguous
pi0 = to_c(pi0)
Ps = to_c(Ps)
ll = to_c(ll)
forward_pass(pi0, Ps, ll, alphas)
return logsumexp(alphas[-1])
def log_transition_matrices(self, data, input, mask, tag):
# Pass the data and inputs through the neural network
x = np.hstack((data[:-1], input[1:]))
for W, b in zip(self.weights, self.biases):
y = np.dot(x, W) + b
x = self.nonlinearity(y)
# Add the baseline transition biases
log_Ps = self.log_Ps[None, :, :] + y[:, None, :]
# Normalize
return log_Ps - logsumexp(log_Ps, axis=2, keepdims=True)
def nll_GLM_GanmorCalciumAR1(w, X, Y, hyperparams, nlfun, S=10):
"""
Negative log-likelihood for a GLM with Ganmor AR1 mixture model for calcium imaging data.
Input:
w: [D x 1] vector of GLM regression weights
X: [T x D] design matrix
Y: [T x 1] calcium fluorescence observations
hyperparams: [3 x 1] model hyperparameters: log tau, log alpha, log Gaussian variance
nlfun: [func] function handle for nonlinearity
S: [scalar] number of spikes to marginalize
return_hess: [bool] flag for returning Hessian
Output:
negative log-likelihood, gradient, and Hessian
"""
# unpack hyperparams
ar_coefs, log_alpha, log_sig2 = hyperparams
alpha = np.exp(log_alpha)
sig2 = np.exp(log_sig2)
p = ar_coefs.shape[0] # AR(p)
# compute AR(p) diffs
Ydecay = np.zeros_like(Y)
Y = np.pad(Y, (p, 0)) # pad Y by p time bins
for i, ai in enumerate(ar_coefs):
Ydecay = Ydecay + ai * Y[p - 1 - i:-1 - i]
# Ydecay2 = ar_coefs[0] * Y[1:-1]
# Ydecay2 = Ydecay2 + ar_coefs[1] * Y[:-2]
# print(np.linalg.norm(Ydecay - Ydecay2))
# import ipdb; ipdb.set_trace()
Ydff = (Y[p:] - Ydecay) / alpha
# compute grid of spike counts
ygrid = np.arange(0, S + 1)
# Gaussian log-likelihood terms
log_gauss_grid = -0.5 * (Ydff[:, None] - ygrid[None, :])**2 / (
sig2 / alpha**2) - 0.5 * np.log(2.0 * np.pi * sig2)
Xproj = X @ w
poissConst = gammaln(ygrid + 1)
# compute neglogli, gradient, and (optionally) Hessian
f, logf, df, ddf = nlfun(Xproj)
logPcounts = logf[:, None] * ygrid[None, :] - f[:,
None] - poissConst[None, :]
# compute log-likelihood for each time bin
logjoint = log_gauss_grid + logPcounts
logli = logsumexp(logjoint, axis=1) # log likelihood for each time bin
negL = -np.sum(logli) # negative log likelihood
return negL
def log_transition_matrices(self, data, input, mask, tag):
def bound_func(t, a, ap, lamb, k):
return a - (1 - np.exp(-(t / lamb)**k)) * (0.0 * a + np.exp(ap))
T, D = data.shape
# Previous state effect
log_Ps = np.tile(self.log_Ps[None, :, :], (T - 1, 1, 1))
# Input effect
boundary_input = 1. - bound_func(input[1:], 1.0, self.ap, self.lamb, 2)
log_Ps = log_Ps + np.dot(boundary_input, self.Ws.T)[:, None, :]
# Past observations effect
log_Ps = log_Ps + np.dot(data[:-1], self.Rs.T)[:, None, :]
return log_Ps - logsumexp(log_Ps, axis=2, keepdims=True)
def hmm_filter(log_pi0, log_Ps, ll):
T, K = ll.shape
# Make sure everything is C contiguous
log_pi0 = to_c(log_pi0)
log_Ps = to_c(log_Ps)
ll = to_c(ll)
# Forward pass gets the predicted state at time t given
# observations up to and including those from time t
alphas = np.zeros((T, K))
forward_pass(log_pi0, log_Ps, ll, alphas)
# Predict forward with the transition matrix
pz_tt = np.exp(alphas - logsumexp(alphas, axis=1, keepdims=True))
pz_tp1t = np.matmul(pz_tt[:-1, None, :], np.exp(log_Ps))[:, 0, :]
# Include the initial state distribution
pz_tp1t = np.row_stack((np.exp(log_pi0 - logsumexp(log_pi0)), pz_tp1t))
assert np.allclose(np.sum(pz_tp1t, axis=1), 1.0)
return pz_tp1t
