def __init__(self, K, pi=None, psi=None, mu=None, Sigma=None):
        """
        Create a PGMultinomial distribution with mean and covariance for psi.

        :param K:       Dimensionality of the multinomial distribution
        :param pi:      Multinomial probability vector (must sum to 1)
        :param psi:     Transformed multinomial probability vector
        :param mu:      Mean of \psi
        :param Sigma:   Covariance of \psi
        """
        assert isinstance(K, int) and K >= 2, "K must be an integer >= 2"
        self.K = K

        if all(param is None for param in (pi,psi,mu,Sigma)):
            mu, sigma = compute_psi_cmoments(np.ones(K))
            Sigma = np.diag(sigma)

        if pi is not None:
            if not (isinstance(pi, np.ndarray) and  pi.shape == (K,)
                    and np.isclose(pi.sum(), 1.0)):
                raise ValueError("Pi must be a normalized length-K vector")
            self.pi = pi

        if psi is not None:
            if not (isinstance(psi, np.ndarray) and psi.shape == (K-1,)):
                raise ValueError("Psi must be a (K-1) vector of reals")
            self.psi = psi

        if mu is not None and Sigma is not None:
            if not (isinstance(mu, np.ndarray) and mu.shape == (K-1,)):
                raise ValueError("Mu must be a (K-1) vector")
            if not (isinstance(Sigma, np.ndarray) and Sigma.shape == ((K-1), (K-1))):
                raise ValueError("Sigma must be a K-1 Covariance matrix")
            self.mu = mu
            self.Sigma = Sigma

            # If psi and pi have not been given, sample from the prior
            if psi is None and pi is None:
                self.psi = np.random.multivariate_normal(self.mu, self.Sigma)

        # Initialize Polya-gamma augmentation variables
        self.ppgs = initialize_polya_gamma_samplers()
        self.omega = np.ones(self.K-1)
dataset/ETHPy150Open HIPS/pgmult/pgmult/distributions.py/PGMultinomial.__init__/original