def _fit_once(): # fit a model if rank is not None: arhmm = ssm.HMM(Kmax, rank, observations="ar") else: arhmm = ssm.HMM(Kmax, N, observations="ar") lls = arhmm.fit(Xp, num_iters=num_iters) return arhmm, lls
def test_hmm_likelihood_perf(T=10000, K=50, D=20): # Create a true HMM A = npr.rand(K, K) A /= A.sum(axis=1, keepdims=True) A = 0.75 * np.eye(K) + 0.25 * A C = npr.randn(K, D) sigma = 0.01 # Sample from the true HMM z = np.zeros(T, dtype=int) y = np.zeros((T, D)) for t in range(T): if t > 0: z[t] = np.random.choice(K, p=A[z[t - 1]]) y[t] = C[z[t]] + np.sqrt(sigma) * npr.randn(D) # Compare to pyhsmm answer from pyhsmm.models import HMM as OldHMM from pybasicbayes.distributions import Gaussian oldhmm = OldHMM( [Gaussian(mu=C[k], sigma=sigma * np.eye(D)) for k in range(K)], trans_matrix=A, init_state_distn="uniform") states = oldhmm.add_data(y) tic = time() true_lkhd = states.log_likelihood() pyhsmm_dt = time() - tic print("PyHSMM: ", pyhsmm_dt, "sec. Val: ", true_lkhd) # Make an HMM with these parameters hmm = ssm.HMM(K, D, observations="gaussian") hmm.transitions.log_Ps = np.log(A) hmm.observations.mus = C hmm.observations._sqrt_Sigmas = np.sqrt(sigma) * np.array( [np.eye(D) for k in range(K)]) tic = time() test_lkhd = hmm.log_probability(y) smm_dt = time() - tic print("SMM HMM: ", smm_dt, "sec. Val: ", test_lkhd) # Make an ARHMM with these parameters arhmm = ssm.HMM(K, D, observations="ar") tic = time() arhmm.log_probability(y) arhmm_dt = time() - tic print("SSM ARHMM: ", arhmm_dt, "sec.") # Make an ARHMM with these parameters arhmm = ssm.HMM(K, D, observations="ar") tic = time() arhmm.expected_states(y) arhmm_dt = time() - tic print("SSM ARHMM Expectations: ", arhmm_dt, "sec.")
def test_viterbi(T=1000, K=20, D=2): # Create a true HMM A = npr.rand(K, K) A /= A.sum(axis=1, keepdims=True) A = 0.75 * np.eye(K) + 0.25 * A C = npr.randn(K, D) sigma = 0.01 # Sample from the true HMM z = np.zeros(T, dtype=int) y = np.zeros((T, D)) for t in range(T): if t > 0: z[t] = np.random.choice(K, p=A[z[t-1]]) y[t] = C[z[t]] + np.sqrt(sigma) * npr.randn(D) # Compare to pyhsmm answer from pyhsmm.models import HMM as OldHMM from pyhsmm.basic.distributions import Gaussian oldhmm = OldHMM([Gaussian(mu=C[k], sigma=sigma * np.eye(D)) for k in range(K)], trans_matrix=A, init_state_distn="uniform") oldhmm.add_data(y) states = oldhmm.states_list.pop() states.Viterbi() z_star = states.stateseq # Make an HMM with these parameters hmm = ssm.HMM(K, D, observations="diagonal_gaussian") hmm.transitions.log_Ps = np.log(A) hmm.observations.mus = C hmm.observations.sigmasq = sigma * np.ones((K, D)) z_star2 = hmm.most_likely_states(y) assert np.allclose(z_star, z_star2)
def test_hmm_likelihood(T=1000, K=5, D=2): # Create a true HMM A = npr.rand(K, K) A /= A.sum(axis=1, keepdims=True) A = 0.75 * np.eye(K) + 0.25 * A C = npr.randn(K, D) sigma = 0.01 # Sample from the true HMM z = np.zeros(T, dtype=int) y = np.zeros((T, D)) for t in range(T): if t > 0: z[t] = np.random.choice(K, p=A[z[t - 1]]) y[t] = C[z[t]] + np.sqrt(sigma) * npr.randn(D) # Compare to pyhsmm answer from pyhsmm.models import HMM as OldHMM from pybasicbayes.distributions import Gaussian oldhmm = OldHMM( [Gaussian(mu=C[k], sigma=sigma * np.eye(D)) for k in range(K)], trans_matrix=A, init_state_distn="uniform") true_lkhd = oldhmm.log_likelihood(y) # Make an HMM with these parameters hmm = ssm.HMM(K, D, observations="diagonal_gaussian") hmm.transitions.log_Ps = np.log(A) hmm.observations.mus = C hmm.observations.sigmasq = sigma * np.ones((K, D)) test_lkhd = hmm.log_probability(y) assert np.allclose(true_lkhd, test_lkhd)
def _fit_hmm(K, X_train, y_train, X_test, y_test, transitions, prior_weight, true_hmm): """ Fits HMM with given parameters 20 times and returns all of them in a list of dictionaries """ X_train = X_train.astype(float) X_test = X_test.astype(float) input_size = X_train.shape[1] hmm = ssm.HMM(K, D, M=input_size, observations="fixedlogistic", transitions=transitions, transition_kwargs={"log_Ps": true_hmm.log_Ps}, observation_kwargs={ "input_size": input_size, "prior_weight": prior_weight, "set_coef": true_hmm.coefs }) lls = hmm.fit(y_train, inputs=X_train, method="em") test_ll = hmm.log_likelihood(y_test, inputs=X_test) result = {} result['hmm'] = hmm result['lls'] = lls result['test_ll'] = test_ll return result
def _fit_hmm(K, X_train, y_train, X_test, y_test, prior_weight, transitions, init_weight): """ Fits HMM with given parameters 20 times and returns all of them in a list of dictionaries """ input_size = X_train.shape[1] hmm = ssm.HMM(K, D, M=input_size, observations="logistic", transitions=transitions, observation_kwargs={ "input_size": input_size, "prior_weight": prior_weight, "init_weight": init_weight }) if init_weight is None: initialize = False else: initialize = True lls = hmm.fit(y_train, inputs=X_train, method="em", initialize=initialize) test_ll = hmm.log_likelihood(y_test, inputs=X_test) result = {} result['hmm'] = hmm result['lls'] = lls result['test_ll'] = test_ll return result
def test_hmm(self): T = 100 # number of time bins K = 5 # number of discrete states D = 2 # dimension of the observations # make an hmm and sample from it hmm = ssm.HMM(K, D, observations="gaussian") z, y = hmm.sample(T) print(z) print(y) #Fitting an HMM is simple. test_hmm = ssm.HMM(K, D, observations="gaussian") test_hmm.fit(y) zhat = test_hmm.most_likely_states(y) print(zhat)
def test_constrained_hmm(T=100, K=3, D=3): hmm = ssm.HMM(K, D, M=0, transitions="constrained", observations="gaussian") z, x = hmm.sample(T) transition_mask = np.array([ [1, 0, 1], [1, 0, 0], [1, 0, 1], ]).astype(bool) init_Ps = np.random.rand(3, 3) init_Ps /= init_Ps.sum(axis=-1, keepdims=True) transition_kwargs = dict( transition_mask=transition_mask ) fit_hmm = ssm.HMM(K, D, M=0, transitions="constrained", observations="gaussian", transition_kwargs=transition_kwargs) fit_hmm.fit(x) learned_Ps = fit_hmm.transitions.transition_matrix assert np.all(learned_Ps[~transition_mask] == 0)
def _model_init(self, dataset): state_list = dataset.get_state_list() observation_list = dataset.get_obs_list() K = len(self._state_list) # number of discrete states D = len(self._observation_list) # dimension of the observation self._hmm = ssm.HMM(K, D, observations='bernoulli') # note: # the log classes govern their own data type # therefore set_pi for example transforms matrix to Prob domain itsself # no action is needed assert self._act_data != None and self._act_data != {} init_pi = gen_handcrafted_priors_for_pi(self._act_data, K) self.set_new_pi(init_pi)
def fitting(trial_num_states, obs_dim, model, obs, optimizer, N_iters): """ fits the model parameter (transition matrix/emission matrix) to the data and returns the likelehoods for each iteration :param trial_num_states: integer, number of hidden states :param obs_dim: integer, dimension of the input channels (original 27, reduced: 2-4) :param model: string, distribution model of the data (eg. gaussian) :param obs: 2d array, (time_bins, obs_dim), data :param optimizer: string ('em' or 'sgd' :param N_iters: integer, number of iterations :return: floats, hmm_lls for each iteraton """ hmm = ssm.HMM(trial_num_states, obs_dim, observations=model) hmm_lls = hmm.fit(obs, method=optimizer, num_iters=N_iters, init_method="kmeans") return hmm, hmm_lls
#import numpy as np import numpy.random as npr import matplotlib.pyplot as plt import ssm from ssm.util import one_hot, find_permutation # %% num_states = 2 # number of discrete states K obs_dim = 1 # number of observed dimensions D num_categories = 2 # number of categories for output C #should not be used here for continous output input_dim = 3 # input dimensions M # Make a GLM-HMM true_glmhmm = ssm.HMM(num_states, obs_dim, input_dim, observations="categorical", observation_kwargs=dict(C=num_categories), transitions="inputdriven") #fix to be both driven! # %% replace from here for now #true_glmhmm.observations = GLM_PoissonObservations(num_states, obs_dim, input_dim) true_glmhmm.observations = InputVonMisesObservations(num_states, obs_dim, input_dim) print(true_glmhmm.transitions.Ws.shape) print(true_glmhmm.observations.mus.shape) # %% #input sequence true_glmhmm.transitions.Ws *= 3 num_sess = 10 # number of example sessions num_trials_per_sess = 1000 # number of trials in a session
npr.seed(0) # In[2]: # Set the parameters of the HMM T = 1000 # number of time bins K = 2 # number of discrete states D = 1 # data dimension M = 1 # input dimension C = 3 # number of output types/categories # Make an HMM true_hmm = ssm.HMM(K, D, M, observations="categorical", observation_kwargs=dict(C=C), transitions="inputdriven") # Optionally, turn up the input weights to exaggerate the effect # true_hmm.transitions.Ws *= 3 # Create an exogenous input inpt = np.sin(2 * np.pi * np.arange(T) / 50)[:, None] + 1e-1 * npr.randn(T, M) # Sample some data from the HMM z, y = true_hmm.sample(T, input=inpt) # Compute the true log probability of the data, summing out the discrete states true_lp = true_hmm.log_probability(y, inputs=inpt)
def loop_over_states_and_statistics(max_num_states, num_loops_statistics, N_iters, e_model, obs, verbose): """ Inner loop: fit the model for the even and calc the log likelihood (llh) with that model for the odd and vice versa the metric for finding a good guess for the number of hidden states is then the sum of these two log likelihoods (Lit: Celeux, Durand). The AIC criterion taking into account the Occram's razor principle (take the simples model which can be described with few params) based on the even-odd cross sum llh is also calculated and stored as well as the "normal" llh. The values are normalized. The codes for fitting and llh calculation are in the ssm package of the linderman's lab. The results is stored in a list which are the value of a dictionary ("dict_indicators"), the keys are the number of hidden states outer loop: loop for statistics, do the inner loop n times to get n plots to compare (do the fit from other init values). Each inner loop produces a dictionary (keys: number of hidden states, values: list of llh's) These dictionaries are stored in a list (appended in each loop) :return: a list of dictionaries ("lst_dicts") containing the values of the e/o llh, AIC and llh for each number of hidden state """ lst_dicts = [] lst_number_states = number_states(max_num_states) obs_even, obs_odd, obs_dim = get_even_odds(obs) # insert loop for statistics (gives in the end num_loop of plots) for j in range(num_loops_statistics): dict_indicators = {} # loop over states for i in range(len(lst_number_states)): # fit model for the even and calc the log likeliood with that model(parmas) for the odd and vice versa # the indicator is then the sum of these two log likelihoods (Lit: Celeux, Durand) # open value list for dictionary of indicators lst = [] even_hmm = ssm.HMM(lst_number_states[i], obs_dim, observations=e_model, transitions="standard") # Fit hmm_lps = even_hmm.fit(obs_even, method="em", num_iters=N_iters, verbose=0) # llh_even = even_hmm.log_likelihood(obs_odd) odd_hmm = ssm.HMM(lst_number_states[i], obs_dim, observations=e_model, transitions="standard") # Fit hmm_lps = odd_hmm.fit(obs_odd, method="em", num_iters=N_iters, verbose=0) # llh_odd = odd_hmm.log_likelihood(obs_even) # store sum of both (the criteria, normalized) sum_llh = even_hmm.log_likelihood( obs_odd) + odd_hmm.log_likelihood(obs_even) lst.append(sum_llh / np.size(obs, 0)) # loglikelihood ohne cross validation hmm = ssm.HMM(lst_number_states[i], obs_dim, observations="Gaussian", transitions="standard") # Fit (append the normalised value) hmm_lps = hmm.fit(obs, method="em", num_iters=N_iters, verbose=0) llh = hmm.log_likelihood(obs) lst.append(llh / np.size(obs, 0)) # calculate AIC from cross sum llh (normalized value) lst.append( AIC(lst_number_states[i], obs_dim, sum_llh) / np.size(obs, 0)) # insert the list as value into dict dict_indicators[lst_number_states[i]] = lst lst_dicts.append(dict_indicators) lst_table = nice_table(lst_dicts) print() print('loop {}:'.format(j)) if verbose: print(lst_table[j].T) return lst_dicts, lst_table
def test_sample(T=10, K=4, D=3, M=2): """ Test that we can construct and sample an HMM with or withou, prefixes, noise, and noise. """ transition_names = [ "standard", "sticky", "inputdriven", "recurrent", "recurrent_only", "rbf_recurrent", "nn_recurrent" ] observation_names = [ "gaussian", "diagonal_gaussian", "t", "diagonal_t", "exponential", "bernoulli", "categorical", "poisson", "vonmises", "ar", "no_input_ar", "diagonal_ar", "independent_ar", "robust_ar", "no_input_robust_ar", "diagonal_robust_ar" ] # Sample basic (no prefix, inputs, etc.) for transitions in transition_names: for observations in observation_names: hmm = ssm.HMM(K, D, M=0, transitions=transitions, observations=observations) zsmpl, xsmpl = hmm.sample(T) # Sample with prefix for transitions in transition_names: for observations in observation_names: hmm = ssm.HMM(K, D, M=0, transitions=transitions, observations=observations) zpre, xpre = hmm.sample(3) zsmpl, xsmpl = hmm.sample(T, prefix=(zpre, xpre)) # Sample with inputs for transitions in transition_names: for observations in observation_names: hmm = ssm.HMM(K, D, M=M, transitions=transitions, observations=observations) zpre, xpre = hmm.sample(3, input=npr.randn(3, M)) zsmpl, xsmpl = hmm.sample(T, prefix=(zpre, xpre), input=npr.randn(T, M)) # Sample without noise for transitions in transition_names: for observations in observation_names: hmm = ssm.HMM(K, D, M=M, transitions=transitions, observations=observations) zpre, xpre = hmm.sample(3, input=npr.randn(3, M)) zsmpl, xsmpl = hmm.sample(T, prefix=(zpre, xpre), input=npr.randn(T, M), with_noise=False)
color=colors[1], label="MF" if n == 0 else None) plt.plot(q_struct_y[:, n] + 4 * n, ':', color=colors[2], label="Struct" if n == 0 else None) plt.legend() plt.xlabel("time") # # Fit an HMM to the LDS states # In[13]: N_iters = 50 K = 15 hmm = ssm.HMM(K, D, observations="gaussian") hmm_lls = hmm.fit(x, method="em", num_em_iters=N_iters) z = hmm.most_likely_states(x) # In[14]: plt.plot(hmm_lls, label="EM") plt.xlabel("EM Iteration") plt.ylabel("Log Probability") plt.legend(loc="lower right") # In[15]: # Plot the observation distributions from hips.plotting.colormaps import white_to_color_cmap
def __init__(self, x_true): self.test_hmm = ssm.HMM(K, D, observations="bernoulli") self.test_hmm.fit(x_true) self.K = K
import autograd.numpy.random as npr npr.seed(0) import matplotlib import matplotlib.pyplot as plt import ssm from ssm.util import find_permutation # Set the parameters of the HMM T = 500 # number of time bins K = 5 # number of discrete states D = 2 # number of observed dimensions # Make an HMM with the true parameters true_hmm = ssm.HMM(K, D, observations="diagonal_gaussian") z, y = true_hmm.sample(T) z_test, y_test = true_hmm.sample(T) true_ll = true_hmm.log_probability(y) # Fit models N_sgd_iters = 1000 N_em_iters = 100 # A bunch of observation models that all include the # diagonal Gaussian as a special case. observations = [ "diagonal_gaussian", "gaussian", "diagonal_t", "studentst", "diagonal_ar", "ar", "diagonal_robust_ar", "robust_ar" ]
plt.plot(x[:, 0], x[:, 1]) # In[7]: mog = MixtureOfGaussians(100, D) mog.fit(x) # In[8]: plt.plot(x[:, 0], x[:, 1]) for mu in mog.observations.mus: plt.plot(mu[0], mu[1], 'o') # In[9]: arhmm = ssm.HMM(K=8, D=D, observations="ar") arhmm.fit(x) # In[10]: z_smpl, x_smpl = arhmm.sample(T=3000) plt.plot(x_smpl[:, 0], x_smpl[:, 1]) # In[11]: def sample(T=3000, num_samples=25, num_burnin=100, schedule=None, filename="samples.mp4"):
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # %% load circuit results obs = rr.T.copy() #rr[0,:][:,None] # inpt = stim.T.copy() #stim[0,:][:,None] # # Set the parameters of the HMM time_bins = obs.shape[1] # number of time bins num_states = 3 # number of discrete states obs_dim = obs.shape[1] # data dimension input_dim = stim.shape[0] # input dimension num_categories = 3 # number of output types/categories # %% run driven HMM inference # Now create a new HMM and fit it to the data with EM N_iters = 100 hmm = ssm.HMM(num_states, obs_dim, input_dim, observations="gaussian", #observation_kwargs=dict(C=num_categories), transitions="inputdriven") # Fit hmm_lps = hmm.fit(obs, inputs=inpt, method="em", num_iters=N_iters) # %% # Plot the log probabilities of the true and fit models plt.figure() plt.plot(hmm_lps, label="EM") plt.legend(loc="lower right") plt.xlabel("EM Iteration") plt.xlim(0, N_iters) plt.ylabel("Log Probability") plt.show()
npr.seed(0) import matplotlib import matplotlib.pyplot as plt import ssm from ssm.util import find_permutation # Set the parameters of the HMM T = 500 # number of time bins K = 5 # number of discrete states D = 2 # number of observed dimensions # Make an HMM with the true parameters true_hmm = ssm.HMM(K, D, observations="exponential") z, y = true_hmm.sample(T) z_test, y_test = true_hmm.sample(T) true_ll = true_hmm.log_probability(y) # Fit models N_sgd_iters = 1000 N_em_iters = 100 # A bunch of observation models that all include the # diagonal Gaussian as a special case. observations = ["exponential"] # Fit with both SGD and EM methods = ["sgd", "em"]
# loads training data from mat-file, saves posterior probabilities into ssm_posterior_probs.mat import numpy import ssm from scipy import io import time import os # Build an HMM instance and set parameters #np.random.seed(1) num_states = 20 # number of discrete states obs_dim = 14 # dimensionality of observation #cov="gaussian" #cov="diagonal_gaussian" cov = 'autoregressive' hmm = ssm.HMM(num_states, obs_dim, observations=cov) #load data using loadmat mat = io.loadmat('training_data.mat') #mat=io.loadmat('C:/Users/Kat/Resilio Sync/Prey Capture/state_epoch_clips-06-Jan-2021/training_data.mat') X = mat['X'] #fit hmm to data N_iters = 200 #N_iters=10 hmm_lls = hmm.fit(X, method="em", num_iters=N_iters) Z = hmm.most_likely_states(X) Ps = hmm.expected_states(X) TM = hmm.transitions.transition_matrix run_on = time.asctime(time.localtime(time.time())) run_from = os.getcwd()
import matplotlib.pyplot as plt import ssm from ssm.util import one_hot, find_permutation import statsmodels.api as sm from scipy.stats import bernoulli npr.seed(0) # Set the parameters of the GLM-HMM num_states = 2 # number of discrete states obs_dim = 1 # number of observed dimensions num_categories = 2 # number of categories for output input_dim = 36 # input dimensions # Make a GLM-HMM true_glmhmm = ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs", observation_kwargs=dict(C=num_categories), transitions="standard") # Put some toy parameters for simulations gen_weights = np.array([[[2.6, 1.3, 0.3, 0.2, 0.1, 0.01, 0.01, 0.01, 0.01, 0.01, 0.06, 0.03, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 1, 0.6, 0.2, 0.1, 0.02, 0.8, 0.4, 0.1, 0.05, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.3]], [[2.0, 1.6, 1.0, 0.6, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 2.0, 1.6, 1.0, 0.6, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.6, 0.2, 0.15, 0.15, 0.05, 0.6, 0.2, 0.15, 0.15, 0.05, 0.8, 0.6, 0.3, 0.1, 0.01, 0.2]]])
def _model_init(self, dataset): K = len(self._state_list) # number of discrete states D = len(self._observation_list) # dimension of the observation self._hmm = ssm.HMM(K, D, observations='bernoulli')
training_inpt = list(map(meta_inpts.__getitem__, train_idx)) training_choice = list( map(meta_true_choices.__getitem__, train_idx)) test_inpt = list(map(meta_inpts.__getitem__, test_idx)) test_choices = list(map(meta_true_choices.__getitem__, test_idx)) for i in np.arange(1, 4): num_states = i obs_dim = 1 input_dim = 5 num_categories = 2 i_ll = [] for p in np.arange(20): bandit_glmhmm = ssm.HMM( num_states, obs_dim, input_dim, observations="input_driven_obs", observation_kwargs=dict(C=num_categories), transitions="standard") fit_ll = bandit_glmhmm.fit(training_choice, inputs=training_inpt, method="em", num_iters=N_iters, tolerance=10**-4) i_ll.append( bandit_glmhmm.log_likelihood(test_choices, inputs=test_inpt)) lls.append(np.mean(i_ll)) lls_var.append(np.stddev(i_ll)) log_likelihoods[:, fold] = lls
import numpy import ssm from scipy import io import time import os # Build an HMM instance and set parameters #np.random.seed(1) num_states = 60 # number of discrete states observation_class = 'autoregressive' obs_dim = 14 # dimensionality of observation transitions = 'sticky' kappa = 1E14 AR_lags = 20 hmm = ssm.HMM(num_states, obs_dim, observations=observation_class, observation_kwargs={'lags':AR_lags}, transitions=transitions, transition_kwargs={'kappa': kappa}) print([num_states, kappa, AR_lags]) #load data using loadmat mat=io.loadmat('training_data.mat') #mat=io.loadmat('C:/Users/Kat/Resilio Sync/Prey Capture/state_epoch_clips-06-Jan-2021/training_data.mat') X = mat['X'] #fit hmm to data N_iters=20 #N_iters=10 hmm_lls = hmm.fit(X, method="em", num_iters=N_iters) Z = hmm.most_likely_states(X) Ps = hmm.expected_states(X) TM = hmm.transitions.transition_matrix
def fit_hmm (obs, num_states=2, obs_dim=1, obs_dist = 'gaussian'): # Make an HMM hmm = ssm.HMM(num_states, obs_dim, observations=obs_dist) hmm_lls = hmm.fit(obs, method="em") return hmm, hmm_lls[-1]
trial = conditioned_trials[i] visual_time = trials['visStim_times'][trial] cue_time = trials['cue_times'][trial] feedback_time = trials['feedback_times'][trial] # generate the spike count histograms t0 = visual_time - pre_stim_dt tf = feedback_time + post_resp_dt [dataset, time_bins] = generate_spike_counts(recording_name, brain_region, neuron_min_score, bin_dt, t0, tf) (n_neurons, n_bins) = dataset.shape # Create a hmm model train_data = dataset.astype(int).T model = ssm.HMM(N_states, n_neurons, observations="poisson") hmm_lls = model.fit(train_data, method="em", num_iters=1000) posterior = model.filter(train_data) # states = model.most_likely_states(train_data) plt.figure(n_trials, figsize=[9, 5]) for s in range(N_states): plt.plot(posterior[:, s], label="State %d" % s) plt.suptitle('Posterior probability of latent states') plt.xlabel(f'time bin ({int(bin_dt*1000)} ms)') plt.ylabel('probability') plt.legend() plt.show()
# Set the parameters of the HMM T = X.shape[0] # number of time bins K = 5 # number of discrete states D = X.shape[1] # number of observed dimensions # Fit with both SGD and EM #methods = ["sgd", "em"] methods = ["em"] results = {} for obs in observations: for method in methods: print("Fitting {} HMM with {}".format(obs, method)) model = ssm.HMM(K, D, observations=obs) train_lls = model.fit(X, method=method) #test_ll = model.log_likelihood(y_test) smoothed_X = model.smooth(X) # Permute to match the true states #model.permute(find_permutation(z, model.most_likely_states(y))) smoothed_z = model.most_likely_states(X) results[(obs, method)] = (model, train_lls, smoothed_z, smoothed_X) # Plot the inferred states fig, axs = plt.subplots(len(observations), 1, figsize=(12, 8)) # Plot the inferred states for i, obs in enumerate(observations):
def main(hparams): if not isinstance(hparams, dict): hparams = vars(hparams) # print hparams to console _print_hparams(hparams) # start at random times (so test tube creates separate folders) np.random.seed(random.randint(0, 1000)) time.sleep(np.random.uniform(1)) # create test-tube experiment hparams, sess_ids, exp = create_tt_experiment(hparams) if hparams is None: print('Experiment exists! Aborting fit') return # build data generator data_generator = build_data_generator(hparams, sess_ids) # #################### # ### CREATE MODEL ### # #################### # get all latents in list n_datasets = len(data_generator) print('collecting observations from data generator...', end='') data_key = 'ae_latents' if hparams['model_class'].find('labels') > -1: data_key = 'labels' latents, trial_idxs = get_latent_arrays_by_dtype( data_generator, sess_idxs=list(range(n_datasets)), data_key=data_key) obs_dim = latents['train'][0].shape[1] hparams['total_train_length'] = np.sum([l.shape[0] for l in latents['train']]) # get separated by dataset as well latents_sess = {d: None for d in range(n_datasets)} trial_idxs_sess = {d: None for d in range(n_datasets)} for d in range(n_datasets): latents_sess[d], trial_idxs_sess[d] = get_latent_arrays_by_dtype( data_generator, sess_idxs=d, data_key=data_key) print('done') if hparams['model_class'] == 'arhmm' or hparams['model_class'] == 'hmm': hparams['ae_model_path'] = os.path.join( os.path.dirname(data_generator.datasets[0].paths['ae_latents'])) hparams['ae_model_latents_file'] = data_generator.datasets[0].paths['ae_latents'] # collect model constructor inputs if hparams['noise_type'] == 'gaussian': if hparams['n_arhmm_lags'] > 0: if hparams['model_class'][:5] != 'arhmm': # 'arhmm' or 'arhmm-labels' raise ValueError('Must specify model_class as arhmm when using AR lags') obs_type = 'ar' else: if hparams['model_class'][:3] != 'hmm': # 'hmm' or 'hmm-labels' raise ValueError('Must specify model_class as hmm when using 0 AR lags') obs_type = 'gaussian' elif hparams['noise_type'] == 'studentst': if hparams['n_arhmm_lags'] > 0: if hparams['model_class'][:5] != 'arhmm': # 'arhmm' or 'arhmm-labels' raise ValueError('Must specify model_class as arhmm when using AR lags') obs_type = 'robust_ar' else: if hparams['model_class'][:3] != 'hmm': # 'hmm' or 'hmm-labels' raise ValueError('Must specify model_class as hmm when using 0 AR lags') obs_type = 'studentst' else: raise ValueError('%s is not a valid noise type' % hparams['noise_type']) if hparams['n_arhmm_lags'] > 0: obs_kwargs = {'lags': hparams['n_arhmm_lags']} obs_init_kwargs = {'localize': True} else: obs_kwargs = None obs_init_kwargs = {} if hparams['kappa'] == 0: transitions = 'stationary' transition_kwargs = None else: transitions = 'sticky' transition_kwargs = {'kappa': hparams['kappa']} print('constructing model...', end='') np.random.seed(hparams['rng_seed_model']) hmm = ssm.HMM( hparams['n_arhmm_states'], obs_dim, observations=obs_type, observation_kwargs=obs_kwargs, transitions=transitions, transition_kwargs=transition_kwargs) hmm.initialize(latents['train']) hmm.observations.initialize(latents['train'], **obs_init_kwargs) # save out hparams as csv and dict hparams['training_completed'] = False export_hparams(hparams, exp) print('done') # #################### # ### TRAIN MODEL ### # #################### # TODO: move fitting into own function # TODO: adopt early stopping strategy from ssm # precompute normalizers n_datapoints = {} n_datapoints_sess = {} for dtype in {'train', 'val', 'test'}: n_datapoints[dtype] = np.vstack(latents[dtype]).size n_datapoints_sess[dtype] = {} for d in range(n_datasets): n_datapoints_sess[dtype][d] = np.vstack(latents_sess[d][dtype]).size for epoch in range(hparams['n_iters'] + 1): # Note: the 0th epoch has no training (randomly initialized model is evaluated) so we cycle # through `n_iters` training epochs print('epoch %03i/%03i' % (epoch, hparams['n_iters'])) if epoch > 0: hmm.fit(latents['train'], method='em', num_iters=1, initialize=False) # export aggregated metrics on train/val data tr_ll = hmm.log_likelihood(latents['train']) / n_datapoints['train'] val_ll = hmm.log_likelihood(latents['val']) / n_datapoints['val'] exp.log({ 'epoch': epoch, 'dataset': -1, 'tr_loss': tr_ll, 'val_loss': val_ll, 'trial': -1}) # export individual session metrics on train/val data for d in range(data_generator.n_datasets): tr_ll = hmm.log_likelihood(latents_sess[d]['train']) / n_datapoints_sess['train'][d] val_ll = hmm.log_likelihood(latents_sess[d]['val']) / n_datapoints_sess['val'][d] exp.log({ 'epoch': epoch, 'dataset': d, 'tr_loss': tr_ll, 'val_loss': val_ll, 'trial': -1}) # export individual session metrics on test data for d in range(n_datasets): for i, b in enumerate(trial_idxs_sess[d]['test']): n = latents_sess[d]['test'][i].size test_ll = hmm.log_likelihood(latents_sess[d]['test'][i]) / n exp.log({'epoch': epoch, 'dataset': d, 'test_loss': test_ll, 'trial': b}) exp.save() # reconfigure model/states by usage zs = [hmm.most_likely_states(x) for x in latents['train']] usage = np.bincount(np.concatenate(zs), minlength=hmm.K) perm = np.argsort(usage)[::-1] hmm.permute(perm) # save model filepath = os.path.join(hparams['expt_dir'], 'version_%i' % exp.version, 'best_val_model.pt') with open(filepath, 'wb') as f: pickle.dump(hmm, f) # ###################### # ### EVALUATE ARHMM ### # ###################### # export states if hparams['export_states']: export_states(hparams, data_generator, hmm) # export training plots if hparams['export_train_plots']: print('creating training plots...', end='') version_dir = os.path.join(hparams['expt_dir'], 'version_%i' % hparams['version']) save_file = os.path.join(version_dir, 'loss_training') export_train_plots(hparams, 'train', loss_type='ll', save_file=save_file) save_file = os.path.join(version_dir, 'loss_validation') export_train_plots(hparams, 'val', loss_type='ll', save_file=save_file) print('done') # update hparams upon successful training hparams['training_completed'] = True export_hparams(hparams, exp) # get rid of unneeded logging info _clean_tt_dir(hparams)
p = np.array([np.real(v)**2 + np.imag(v)**2 for v in f]) pi = np.fft.ifft(p) return np.real(pi)[:int(x.size / 2)] / np.sum(xp**2) plt.figure() plt.plot(autocorrelation(rt[1, :])[:5000]) # %% Dwell-time calculation! import ssm # %% based on HMM obs = rt[:4, :].T obs_dims = obs.shape[1] N_iters = 50 num_states = 2 #assuming transition between two patterns hmm = ssm.HMM(num_states, obs_dims, num_iters='gaussian') hmm_lls = hmm.fit(obs, method='em', num_iters=N_iters) # %% plt.figure() plt.plot(hmm_lls) most_lls = hmm.most_likely_states(obs) plt.figure() plt.subplot(211) plt.plot(most_lls) plt.xlim([0, len(most_lls)]) plt.subplot(212) plt.imshow(obs.T, aspect='auto') learned_transition_mat = hmm.transitions.transition_matrix print(learned_transition_mat)