def categorical_logpdf(data, logits, mask=None):
"""
Compute the log probability density of a categorical distribution.
This will broadcast as long as data and logits have the same
(or at least compatible) leading dimensions.
Parameters
----------
data : array_like (..., D) int (0 <= data < C)
The points at which to evaluate the log density
lambdas : array_like (..., D, C)
The logits of the categorical distribution(s) with C classes
mask : array_like (..., D) bool
Optional mask indicating which entries in the data are observed
Returns
-------
lps : array_like (...,)
Log probabilities under the categorical distribution(s).
"""
D = data.shape[-1]
C = logits.shape[-1]
assert data.dtype in (int, np.int8, np.int16, np.int32, np.int64)
assert np.all((data >= 0) & (data < C))
assert logits.shape[-2] == D
# Check mask
mask = mask if mask is not None else np.ones_like(data, dtype=bool)
assert mask.shape == data.shape
logits = logits - logsumexp(logits, axis=-1, keepdims=True) # (..., D, C)
x = one_hot(data, C) # (..., D, C)
lls = np.sum(x * logits, axis=-1) # (..., D)
return np.sum(lls * mask, axis=-1) # (...,)
def initialize(self,
datas,
inputs=None,
masks=None,
tags=None,
init_method="random"):
Ts = [data.shape[0] for data in datas]
# Get initial discrete states
if init_method.lower() == 'kmeans':
# KMeans clustering
from sklearn.cluster import KMeans
km = KMeans(self.K)
km.fit(np.vstack(datas))
zs = np.split(km.labels_, np.cumsum(Ts)[:-1])
elif init_method.lower() == 'random':
# Random assignment
zs = [npr.choice(self.K, size=T) for T in Ts]
else:
raise Exception(
'Not an accepted initialization type: {}'.format(init_method))
# Make a one-hot encoding of z and treat it as HMM expectations
Ezs = [one_hot(z, self.K) for z in zs]
expectations = [(Ez, None, None) for Ez in Ezs]
# Set the variances all at once to use the setter
self.m_step(expectations, datas, inputs, masks, tags)
def m_step(self, expectations, datas, inputs, masks, tags, **kwargs):
x = np.concatenate(datas)
weights = np.concatenate([Ez for Ez, _, _ in expectations])
for k in range(self.K):
# compute weighted histogram of the class assignments
xoh = one_hot(x, self.C) # T x D x C
ps = np.average(xoh, axis=0, weights=weights[:, k]) + 1e-3 # D x C
ps /= np.sum(ps, axis=-1, keepdims=True)
self.logits[k] = np.log(ps)
def log_likelihoods(self, data, input, mask, tag):
assert (data.dtype == int or data.dtype == bool)
assert data.ndim == 2 and data.shape[1] == self.D
assert data.min() >= 0 and data.max() < self.C
logits = self.logits - logsumexp(self.logits, axis=2, keepdims=True) # K x D x C
x = one_hot(data, self.C) # T x D x C
lls = np.sum(x[:, None, :, :] * logits, axis=3) # T x K x D
mask = np.ones_like(data, dtype=bool) if mask is None else mask # T x D
return np.sum(lls * mask[:, None, :], axis=2) # T x K
def m_step(self, expectations, datas, inputs, masks, tags, optimizer="adam", num_iters=10, **kwargs):
"""
Fit a logistic regression for the transitions.
Technically, this is a stochastic M-step since the states
are sampled from their posterior marginals.
"""
from sklearn.linear_model import LogisticRegression
K, M, D = self.K, self.M, self.D
zps, zns = [], []
for Ez, _ in expectations:
z = np.array([np.random.choice(K, p=p) for p in Ez])
zps.append(z[:-1])
zns.append(z[1:])
X = np.vstack([np.hstack((one_hot(zp, K), input[1:], data[:-1]))
for zp, input, data in zip(zps, inputs, datas)])
y = np.concatenate(zns)
# Determine the number of states used
used = np.unique(y)
K_used = len(used)
unused = np.setdiff1d(np.arange(K), used)
# Reset parameters before filling in
self.log_Ps = np.zeros((K, K))
self.Ws = np.zeros((K, M))
self.Rs = np.zeros((K, D))
if K_used == 1:
warn("RecurrentTransitions: Only using 1 state in expectation. "
"M-step cannot proceed. Resetting transition parameters.")
return
# Fit the logistic regression
lr = LogisticRegression(fit_intercept=False, multi_class="multinomial", solver="sag")
lr.fit(X, y)
# Extract the coefficients
assert lr.coef_.shape[0] == (K_used if K_used > 2 else 1)
log_P = lr.coef_[:, :K]
W = lr.coef_[:, K:K+M]
R = lr.coef_[:, K+M:]
if K_used == 2:
# lr thought there were only two classes
self.log_Ps[:,used[1]] = lr.coef_[0, :K]
self.Ws[used[1]] = lr.coef_[0,K:K+M]
self.Rs[used[1]] = lr.coef_[0,K+M:]
else:
self.log_Ps[:, used] = log_P.T
self.Ws[used] = W
self.Rs[used] = R
def simulate_ramping(beta=np.linspace(-0.02, 0.02, 5),
w2=3e-3,
x0=0.5,
C=40,
T=100,
bin_size=0.01):
NC = 5 # number of trial types
cohs = np.arange(NC)
trial_cohs = np.repeat(cohs, int(T / NC))
tr_lengths = np.random.randint(50, size=(T)) + 50
us = []
xs = []
zs = []
ys = []
for t in range(T):
tr_coh = trial_cohs[t]
betac = beta[tr_coh]
tr_length = tr_lengths[t]
x = np.zeros(tr_length)
z = np.zeros(tr_length)
x[0] = x0 + np.sqrt(w2) * npr.randn()
z[0] = 0
for i in np.arange(1, tr_length):
if x[i - 1] >= 1.0:
x[i] = 1.0
z[i] = 1
else:
x[i] = np.min(
(1.0, x[i - 1] + betac + np.sqrt(w2) * npr.randn()))
if x[i] >= 1.0:
z[i] = 1
else:
z[i] = 0
y = npr.poisson(np.log1p(np.exp(C * x)) * bin_size)
u = np.tile(one_hot(tr_coh, 5), (tr_length, 1))
us.append(u)
xs.append(x.reshape((tr_length, 1)))
zs.append(z.reshape((tr_length, 1)))
ys.append(y.reshape((tr_length, 1)))
return ys, xs, zs, us, tr_lengths, trial_cohs
latent_ramp.emissions.Cs[0] = 40.0 + 3.0 * npr.randn(N,1) # Simulate data numTrials = 100 ys = [] xs = [] zs = [] us = [] cohs = [] # sample from model for tr in range(numTrials): coh = npr.randint(M) u_tr = one_hot(coh,M) tr_length = npr.randint(50)+50 u = u_tr * np.ones((tr_length,1)) T = np.shape(u)[0] z, x, y = latent_ramp.sample(T, input=u) zs.append(z) xs.append(x) ys.append(y) us.append(u) cohs.append(coh) cohs = np.array(cohs) def initialize_ramp(ys,cohs, bin_size):
plt.subplot(311)
plt.plot(inpt)
plt.xticks([])
plt.xlim(0, T)
plt.ylabel("input")
plt.subplot(312)
plt.imshow(z[None, :], aspect="auto")
plt.xticks([])
plt.xlim(0, T)
plt.ylabel("discrete\nstate")
plt.yticks([])
plt.subplot(313)
plt.imshow(one_hot(y[:, 0], C).T, aspect="auto")
plt.xlim(0, T)
plt.ylabel("observation")
plt.yticks(np.arange(C))
# In[4]:
# Now create a new HMM and fit it to the data with EM
N_iters = 50
hmm = ssm.HMM(K,
D,
M,
observations="categorical",
observation_kwargs=dict(C=C),
transitions="inputdriven")
smoothed_z_gauss_test = test_hmm_gauss.most_likely_states(y_ca_test)
plt.figure()
plt.imshow(np.row_stack((z_test, smoothed_z_test, smoothed_z_gauss_test, smoothed_z_ca_test)), aspect="auto")
plt.xlim([0, T_plot])
plt.yticks([0,1,2,3],["true","poisson","gaussian","calcium"])
plt.ylim([3.5,-0.5])
# posterior expectations
poiss_expectations = test_hmm.expected_states(y)[0]
gauss_expectations = test_hmm_gauss.expected_states(y_ca)[0]
ganmor_expectations = test_hmm_ca.expected_states(y_ca)[0]
plt.figure()
plt.subplot(311)
plt.imshow(one_hot(z, K=K).T, aspect="auto", vmin=0.0, vmax=1.0, cmap="Greys")
plt.title("true")
plt.xlim([0, T_plot])
plt.subplot(312)
plt.imshow(gauss_expectations.T, aspect="auto", vmin=0.0, vmax=1.0, cmap="Greys")
plt.xlim([0, T_plot])
plt.title("gaussian")
plt.subplot(313)
plt.imshow(ganmor_expectations.T, aspect="auto", vmin=0.0, vmax=1.0, cmap="Greys")
plt.xlim([0, T_plot])
plt.title("ganmor")
plt.figure()
plt.subplot(221)
plt.imshow(true_hmm.transitions.transition_matrix, aspect="auto", vmin=0.0, vmax=1.0)
plt.title("true")
