Simple implementation for a Gaussian mixture distribution

gmd.py

import numpy as np
import scipy.stats as st

class GaussianMixture(object):
    def __init__(self, mu, sigma, w):
        self.mu = np.asarray(mu)
        self.sigma = np.asarray(sigma)
        
        self.w = np.asarray(w)
        self.w /= self.w.sum()
        self.n_comp = len(self.w)
        
        self._rv = st.norm(loc=self.mu, scale=self.sigma)
        
    def pdf(self, X):
        # Calculate PDF as a w-weighted sum of component PDFs.
        # Because of how st.norm works with multiple values, we need to 
        # pass it a matrix whose columns are n_comp copies of X
        Xx = np.repeat(np.atleast_2d(X), self.n_comp, axis=0).T
        
        # Matrix whose i-th column is the PDF of the i-th component evaluated
        # on the grid X
        P = self._rv.pdf(Xx)
        
        # This einsum expression multiplies each row of P by w and then sums the
        # columns of the result
        return np.einsum(
            "i,ji->j",
            self.w,
            P
        )
    
    def rvs(self, N=1, return_components=False):
        # Choose component indices
        I = np.random.choice(self.n_comp, size=N, replace=True, p=self.w)
        
        # Generate a vector of means and variances for each sample corresponding
        # to the chosen components
        mu_N = self.mu[I]
        sigma_N = self.sigma[I]
        
        # Sample
        rvs = st.norm.rvs(loc=mu_N, scale=sigma_N)
        
        if return_components:
            return rvs, I # Include a list of the chosen components
        else:
            return rvs 
    
    def __getitem__(self, i):
        # Get a scipy rv_frozen object corresponding to the i-th component distribution.
        # This could be vectorized but I don't need to.
        return st.norm(loc=self.mu[i], scale=self.sigma[i])