def resample_omega(self, augmented_data_list):
        """
        Resample omega from its conditional Polya-gamma distribution
        :return:
        """
        K = self.K
        for data in augmented_data_list:
            x = data["x"]
            T = data["T"]

            # TODO: Fix this hack
            if "z" in data:
                z = data["z"]
            elif "states" in data:
                z = data["states"].stateseq
            else:
                raise Exception("Could not find latent states in augmented data!")

            psi = z.dot(self.C.T) + self.mu[None, :]
            N = N_vec(x).astype(np.float)
            tmp_omg = np.zeros(N.size)
            ppg.pgdrawvpar(self.ppgs, N.ravel(), psi.ravel(), tmp_omg)
            data["omega"] = tmp_omg.reshape((T, self.K-1))

            # Clip out zeros
            data["omega"] = np.clip(data["omega"], 1e-8,np.inf)
dataset/ETHPy150Open HIPS/pgmult/pgmult/distributions.py/PGMultinomialRegression.resample_omega/original