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