jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def kfoldCrossValidationNBG(D, L):
    # Same procedure of naiveBayesGaussianClassifier but with leave-one-out approach:
    # number of folds is equal to the number of samples, we iterate over all samples
    # and each time we use one of them as evaluation set and the remaining 149
    # as training set.
    # Then we compute the predicted lable and count how many times it is correct.
    numberOfCorrectPredictions = 0
    for i in range(D.shape[1]):  # i=0...149
        # Each time we need to delete the i-th column (sample) from the training set
        # while in the evaluation set we should keep only the i-th column.
        # vcol is necessary to exploit broadcasting
        (DTR, LTR), (DEV, LEV) = (np.delete(D, i, 1),
                                  np.delete(L, i)), (vcol(D[:, i], 4), L[i])
        (mu0, sigma0), (mu1, sigma1), (mu2, sigma2) = computeMLestimates(DTR, LTR)
        (sigma0, sigma1, sigma2) = (sigma0*np.identity(sigma0.shape[0]), sigma1*np.identity(
            sigma1.shape[0]), sigma2*np.identity(sigma2.shape[0]))
        logS = computeScoreMatrix(DEV, mu0, sigma0, mu1,
                                  sigma1, mu2, sigma2, logpdf_GAU_ND)
        priorLogProbabilities = vcol(
            np.array([np.log(1/3), np.log(1/3), np.log(1/3)]), 3)
        logSJoint = logS+priorLogProbabilities  # 3x1
        marginalLogDensities = vrow(
            sc.special.logsumexp(logSJoint, axis=0), 1)  # 1x1
        # We need to keep it as 1x1 to exploit broadcasting
        logSPost = logSJoint-marginalLogDensities
        predictedLabels = logSPost.argmax(axis=0)
        # Add the result of prediction to the proper variable
        numberOfCorrectPredictions += np.array(predictedLabels == LEV).sum()
    # Now we can compute percentage values for accuracy and error rate.
    accuracy = numberOfCorrectPredictions/L.size * \
        100  # over all samples, that are 150
    errorRate = 100-accuracy
    print("Accuracy of the Naive Bayes Gaussian classifier when working with log-densities and "
          "using K-fold cross validation with leave-one-out approach: %.2f %%" % (accuracy))
    print("Error rate: %.2f %%" % (errorRate))
    return
//  [python]
def kfoldCrossValidationTCG(D, L):
    numberOfCorrectPredictions = 0
    for i in range(D.shape[1]):
        (DTR, LTR), (DEV, LEV) = (np.delete(D, i, 1),
                                  np.delete(L, i)), (vcol(D[:, i], 4), L[i])
        (mu0, sigma0), (mu1, sigma1), (mu2, sigma2) = computeMLestimates(DTR, LTR)
        sigma = (1/DTR.shape[1])*((LTR == 0).sum()*sigma0 +
                                  (LTR == 1).sum()*sigma1+(LTR == 2).sum()*sigma2)
        logS = computeScoreMatrix(DEV, mu0, sigma, mu1,
                                  sigma, mu2, sigma, logpdf_GAU_ND)
        priorLogProbabilities = vcol(
            np.array([np.log(1/3), np.log(1/3), np.log(1/3)]), 3)
        logSJoint = logS+priorLogProbabilities  # 3x1
        marginalLogDensities = vrow(
            np.log(np.sum(np.exp(logSJoint), axis=0)), 1)  
        logSPost = logSJoint-marginalLogDensities
        predictedLabels = logSPost.argmax(axis=0)
        numberOfCorrectPredictions += np.array(predictedLabels == LEV).sum()
    accuracy = numberOfCorrectPredictions/L.size * 100
    errorRate = 100-accuracy
    print("Accuracy of the Tied Covariance Gaussian classifier when working with log-densities and "
          "using K-fold cross validation with leave-one-out approach: %.2f %%" % (accuracy))
    print("Error rate: %.2f %%" % (errorRate))
    return
//  [python]
def kfoldCrossValidationTCG(D, L):
    # Same procedure of tiedCovarianceGaussianClassifier but with leave-one-out approach:
    # number of folds is equal to the number of samples, we iterate over all samples
    # and each time we use one of them as evaluation set and the remaining 149
    # as training set.
    # Then we compute the predicted lable and count how many times it is correct.
    numberOfCorrectPredictions = 0
    for i in range(D.shape[1]):  # i=0...149
        # Each time we need to delete the i-th column (sample) from the training set
        # while in the evaluation set we should keep only the i-th column.
        # vcol is necessary to exploit broadcasting
        (DTR, LTR), (DEV, LEV) = (np.delete(D, i, 1),
                                  np.delete(L, i)), (vcol(D[:, i], 4), L[i])
        (mu0, sigma0), (mu1, sigma1), (mu2, sigma2) = computeMLestimates(DTR, LTR)
        sigma = (1/DTR.shape[1])*((LTR == 0).sum()*sigma0 +
                                  (LTR == 1).sum()*sigma1+(LTR == 2).sum()*sigma2)
        logS = computeScoreMatrix(DEV, mu0, sigma, mu1,
                                  sigma, mu2, sigma, logpdf_GAU_ND)
        priorLogProbabilities = vcol(
            np.array([np.log(1/3), np.log(1/3), np.log(1/3)]), 3)
        logSJoint = logS+priorLogProbabilities  # 3x1
        marginalLogDensities = vrow(
            sc.special.logsumexp(logSJoint, axis=0), 1)  # 1x1
        # We need to keep it as 1x1 to exploit broadcasting
        logSPost = logSJoint-marginalLogDensities
        predictedLabels = logSPost.argmax(axis=0)
        # Add the result of prediction to the proper variable
        numberOfCorrectPredictions += np.array(predictedLabels == LEV).sum()
    # Now we can compute percentage values for accuracy and error rate.
    accuracy = numberOfCorrectPredictions/L.size * \
        100  # over all samples, that are 150
    errorRate = 100-accuracy
    print("Accuracy of the Tied Covariance Gaussian classifier when working with log-densities and "
          "using K-fold cross validation with leave-one-out approach: %.2f %%" % (accuracy))
    print("Error rate: %.2f %%" % (errorRate))
    return
//  [python]
def _softmax(x):
    lse_x = np.logaddexp.reduce(x, axis=-1, keepdims=True)
    return np.exp(x - lse_x)
//  [python]
def _softmax(x):
    lse_x = scipy.special.logsumexp(x, axis=-1, keepdims=True)
    return np.exp(x - lse_x)
//  [python]
def meanfield_predictive_weights(self, x, dist='gaussian'):
        log_posterior_predictive = self.basis.log_posterior_predictive_gaussian(x)\
            if dist == 'gaussian' else self.basis.log_posterior_predictive_studentt(x)

        log_gating = np.log(self.gating.posterior.mean())
        log_weight = np.expand_dims(log_gating, axis=1) + log_posterior_predictive
        weights = np.exp(log_weight - logsumexp(log_weight, axis=0, keepdims=True))

        return weights
//  [python]
def meanfield_predictive_activation(self, x, dist='gaussian'):
        # for plotting basis functions
        x = np.reshape(x, (-1, self.input_dim))
        xx = self.input_transform.transform(x) if self.scale else x

        log_basis = self.basis.log_posterior_predictive_gaussian(xx)\
            if dist == 'gaussian' else self.basis.log_posterior_predictive_studentt(xx)

        log_gating = np.log(self.gating.posterior.mean())
        log_activations = np.expand_dims(log_gating, axis=1) + log_basis
        activations = np.exp(log_activations - np.max(log_activations, axis=0, keepdims=True)) - np.log(np.sum(np.exp(log_activations - np.max(log_activations, axis=0, keepdims=True)), axis=0, keepdims=True))

        return activations
//  [python]
def meanfield_predictive_activation(self, x, dist='gaussian'):
        # for plotting basis functions
        x = np.reshape(x, (-1, self.input_dim))
        xx = self.input_transform.transform(x) if self.scale else x

        log_basis = self.basis.log_posterior_predictive_gaussian(xx)\
            if dist == 'gaussian' else self.basis.log_posterior_predictive_studentt(xx)

        log_gating = np.log(self.gating.posterior.mean())
        log_activations = np.expand_dims(log_gating, axis=1) + log_basis
        activations = np.exp(log_activations - logsumexp(log_activations, axis=0, keepdims=True))

        return activations
//  [python]
def loglike(self, theta):
        thetas = []
        for pos in self.positions:
            thetas.append(copy.deepcopy(theta))
            thetas[-1][:2] -= pos
        logls = [(log_gaussian_pdf(thetas[i], sigma=self.sigmas[i])
                  + np.log(self.weights[i])) for i in range(len(self.weights))]
        return np.log(np.sum(np.exp(logls)))
//  [python]
def loglike(self, theta):
        thetas = []
        for pos in self.positions:
            thetas.append(copy.deepcopy(theta))
            thetas[-1][:2] -= pos
        logls = [(log_gaussian_pdf(thetas[i], sigma=self.sigmas[i])
                  + np.log(self.weights[i])) for i in range(len(self.weights))]
        return scipy.special.logsumexp(logls)
//  [python]
def log_likelihood(self):
        if not hasattr(self,'_normalizer') or self._normalizer is None:
            scores = self._compute_scores()
            self._normalizer = np.log(np.sum(np.exp(scores[~np.isnan(self.data).any(1)],axis=1))).sum()
        return self._normalizer
//  [python]
def log_likelihood(self):
        if not hasattr(self,'_normalizer') or self._normalizer is None:
            scores = self._compute_scores()
            self._normalizer = logsumexp(scores[~np.isnan(self.data).any(1)],axis=1).sum()
        return self._normalizer
//  [python]
def hmm_expected_states(pi0, Ps, ll):
    T, K = ll.shape

    alphas = np.zeros((T, K))
    forward_pass(pi0, Ps, ll, alphas)
    normalizer = s.logsumexp(alphas[-1])

    betas = np.zeros((T, K))
    backward_pass(Ps, ll, betas)

    # Compute E[z_t] for t = 1, ..., T
    expected_states = alphas + betas
    expected_states -= s.logsumexp(expected_states, axis=1, keepdims=True)
    expected_states = np.exp(expected_states)

    # Compute the log transition matrices.
    # Suppress log(0) warnings as they are expected.
    with np.errstate(divide="ignore"):
        log_Ps = np.log(Ps)


    # Compute E[z_t, z_{t+1}] for t = 1, ..., T-1
    # Note that this is an array of size T*K*K, which can be quite large.
    # To be a bit more frugal with memory, first check if the given log_Ps
    # are TxKxK.  If so, instantiate the full expected joints as well, since
    # we will need them for the M-step.  However, if log_Ps is 1xKxK then we
    # know that the transition matrix is stationary, and all we need for the
    # M-step is the sum of the expected joints.
    stationary = (Ps.shape[0] == 1)
    if not stationary:
        expected_joints = alphas[:-1,:,None] + betas[1:,None,:] + ll[1:,None,:] + log_Ps
        expected_joints -= expected_joints.max((1,2))[:,None, None]
        expected_joints = np.exp(expected_joints)
        expected_joints /= expected_joints.sum((1,2))[:,None,None]

    else:
        # Compute the sum over time axis of the expected joints
        expected_joints = np.zeros((K, K))
        _compute_stationary_expected_joints(alphas, betas, ll, log_Ps[0], expected_joints)
        expected_joints = expected_joints[None, :, :]

    return expected_states, expected_joints, normalizer
//  [python]
def hmm_expected_states(pi0, Ps, ll):
    T, K = ll.shape

    alphas = np.zeros((T, K))
    forward_pass(pi0, Ps, ll, alphas)
    normalizer = logsumexp(alphas[-1])

    betas = np.zeros((T, K))
    backward_pass(Ps, ll, betas)

    # Compute E[z_t] for t = 1, ..., T
    expected_states = alphas + betas
    expected_states -= scsp.logsumexp(expected_states, axis=1, keepdims=True)
    expected_states = np.exp(expected_states)

    # Compute the log transition matrices.
    # Suppress log(0) warnings as they are expected.
    with np.errstate(divide="ignore"):
        log_Ps = np.log(Ps)


    # Compute E[z_t, z_{t+1}] for t = 1, ..., T-1
    # Note that this is an array of size T*K*K, which can be quite large.
    # To be a bit more frugal with memory, first check if the given log_Ps
    # are TxKxK.  If so, instantiate the full expected joints as well, since
    # we will need them for the M-step.  However, if log_Ps is 1xKxK then we
    # know that the transition matrix is stationary, and all we need for the
    # M-step is the sum of the expected joints.
    stationary = (Ps.shape[0] == 1)
    if not stationary:
        expected_joints = alphas[:-1,:,None] + betas[1:,None,:] + ll[1:,None,:] + log_Ps
        expected_joints -= expected_joints.max((1,2))[:,None, None]
        expected_joints = np.exp(expected_joints)
        expected_joints /= expected_joints.sum((1,2))[:,None,None]

    else:
        # Compute the sum over time axis of the expected joints
        expected_joints = np.zeros((K, K))
        _compute_stationary_expected_joints(alphas, betas, ll, log_Ps[0], expected_joints)
        expected_joints = expected_joints[None, :, :]

    return expected_states, expected_joints, normalizer
//  [python]
def _log_softmax(x, axis=-1):
  x_max = np.max(x, axis=axis, keepdims=True)
  x_shifted = x - x_max
  return x_shifted - s.logsumexp(x_shifted, axis=axis, keepdims=True)
//  [python]
def _log_softmax(x, axis=-1):
  x_max = np.max(x, axis=axis, keepdims=True)
  x_shifted = x - x_max
  return x_shifted - sp_special.logsumexp(x_shifted, axis=axis, keepdims=True)
//  [python]
def _log_likelihoods(self,x):
        # NOTE: nans propagate as nans
        x = np.asarray(x)
        K = len(self.components)
        vals = np.empty((x.shape[0],K))
        for idx, c in enumerate(self.components):
            vals[:,idx] = c.log_likelihood(x)
        vals += self.weights.log_likelihood(np.arange(K))
        return logsumexp(vals,axis=1)
//  [python]
def lauricella_2f1_exact_second_way(a,b,c,z):
    
    prefac=-log_beta(a,c-a)
    summands=[]
    signs=[]
    for i in range(c-a):
        signs.append((-1)**i)
        summands.append(lauricella_2f1_exact(a+i,b,a+i+1,z)+log_bin_fac(c-a-1,i)-numpy.log(a+i))
    
    #print "summands ", summands
    #print "signs", signs
    return prefac+scipy.special.logsumexp(summands, b=signs)
//  [python]
def log_lauricella_dumb_integral_single(a,bs, c, zs, nthrows=1000):
    """
    Second known integral lauricella representation in form of a single integral over [0,1]. Only valid if c>a.
    """
    if(a==c):
        ## get out since this is a limiting case and not defined in the integral rep
        return (-bs*numpy.log(1.0-zs)).sum()

    random_vals=numpy.random.uniform(size=nthrows)
    random_vals2=numpy.resize(random_vals, (len(bs),nthrows)).T
    
    log_res=(a-1.0)*numpy.log(random_vals) + (c-a-1.0)* numpy.log(1.0-random_vals) + (  -bs*numpy.log(1.0-random_vals2*zs)).sum( axis=1)- numpy.log(float(nthrows))

    log_normfactor=scipy.special.gammaln(c)-scipy.special.gammaln(a)-scipy.special.gammaln(c-a)

    return numpy.logaddexp.reduce(log_normfactor+log_res)
//  [python]
def log_lauricella_dumb_integral_single(a,bs, c, zs, nthrows=1000):
    """
    Second known integral lauricella representation in form of a single integral over [0,1]. Only valid if c>a.
    """
    if(a==c):
        ## get out since this is a limiting case and not defined in the integral rep
        return (-bs*numpy.log(1.0-zs)).sum()

    random_vals=numpy.random.uniform(size=nthrows)
    random_vals2=numpy.resize(random_vals, (len(bs),nthrows)).T
    
    log_res=(a-1.0)*numpy.log(random_vals) + (c-a-1.0)* numpy.log(1.0-random_vals) + (  -bs*numpy.log(1.0-random_vals2*zs)).sum( axis=1)- numpy.log(float(nthrows))

    log_normfactor=scipy.special.gammaln(c)-scipy.special.gammaln(a)-scipy.special.gammaln(c-a)

    return scipy.special.logsumexp(log_normfactor+log_res)
//  [python]
def _log_softmax(logits):
  log_z = scipy.special.logsumexp(logits)
  return logits - log_z
//  [python]
def meanfield_prediction(self, x, y=None,
                             prediction='average',
                             dist='gaussian',
                             incremental=False,
                             variance='diagonal'):

        x = np.reshape(x, (-1, self.input_dim))

        compute_nlpd = False
        if y is not None:
            y = np.reshape(y, (-1, self.output_dim))
            compute_nlpd = True

        if self.scale:
            xx = self.input_transform.transform(x)
            yy = self.output_transform.transform(y) if y is not None else None
        else:
            xx, yy = x, y if y is not None else None

        weights = self.meanfield_predictive_weights(xx, dist)
        mus, covars = self.meanfield_predictive_moments(xx, dist)

        if prediction == 'mode':
            k = np.argmax(weights, axis=0)
            idx = (k, range(len(k)), ...)
            mu, covar = mus[idx], covars[idx]
        elif prediction == 'average':
            mu, covar = self.mixture_moments(mus, covars, weights)
        else:
            raise NotImplementedError

        nlpd = None
        if compute_nlpd:
            log_pl = self.meanfiled_log_predictive_likelihood(xx, yy)
            log_weights = np.log(weights + eps)
            nlpd = - 1.0 * np.max(log_pl + log_weights, axis=0) + np.log(np.sum(np.exp(log_pl + log_weights - np.max(log_pl + log_weights, axis=0)), axis=0))

        if self.scale:
            mu = self.output_transform.inverse_transform(mu)
            mat = np.diag(np.sqrt(self.output_transform.var_))
            covar = np.einsum('kh,...hj,ji->...ki', mat, covar, mat.T)

        if incremental:
            mu += x[:, :self.output_dim]

        var = np.vstack(list(map(np.diag, covar)))

        if compute_nlpd:
            if variance == 'diagonal':
                return mu, var, np.sqrt(var), nlpd
            else:
                return mu, covar, np.sqrt(var), nlpd
        else:
            if variance == 'diagonal':
                return mu, var, np.sqrt(var)
            else:
                return mu, covar, np.sqrt(var)
//  [python]
def meanfield_prediction(self, x, y=None,
                             prediction='average',
                             dist='gaussian',
                             incremental=False,
                             variance='diagonal'):

        x = np.reshape(x, (-1, self.input_dim))

        compute_nlpd = False
        if y is not None:
            y = np.reshape(y, (-1, self.output_dim))
            compute_nlpd = True

        if self.scale:
            xx = self.input_transform.transform(x)
            yy = self.output_transform.transform(y) if y is not None else None
        else:
            xx, yy = x, y if y is not None else None

        weights = self.meanfield_predictive_weights(xx, dist)
        mus, covars = self.meanfield_predictive_moments(xx, dist)

        if prediction == 'mode':
            k = np.argmax(weights, axis=0)
            idx = (k, range(len(k)), ...)
            mu, covar = mus[idx], covars[idx]
        elif prediction == 'average':
            mu, covar = self.mixture_moments(mus, covars, weights)
        else:
            raise NotImplementedError

        nlpd = None
        if compute_nlpd:
            log_pl = self.meanfiled_log_predictive_likelihood(xx, yy)
            log_weights = np.log(weights + eps)
            nlpd = - 1.0 * logsumexp(log_pl + log_weights, axis=0)

        if self.scale:
            mu = self.output_transform.inverse_transform(mu)
            mat = np.diag(np.sqrt(self.output_transform.var_))
            covar = np.einsum('kh,...hj,ji->...ki', mat, covar, mat.T)

        if incremental:
            mu += x[:, :self.output_dim]

        var = np.vstack(list(map(np.diag, covar)))

        if compute_nlpd:
            if variance == 'diagonal':
                return mu, var, np.sqrt(var), nlpd
            else:
                return mu, covar, np.sqrt(var), nlpd
        else:
            if variance == 'diagonal':
                return mu, var, np.sqrt(var)
            else:
                return mu, covar, np.sqrt(var)
//  [python]
def log_softmax(x, axis):
    return x - s.logsumexp(x, axis=axis, keepdims=True)
//  [python]
def log_softmax(x, axis):
    return x - scipy.special.logsumexp(x, axis=axis, keepdims=True)
//  [python]
def normalize_log_weights(lw):
    """Normalises a matrix of log-weights, row-wise"""

    ref = lw.max(0)
    lsum = np.log(np.exp(lw - ref).sum(0)) + ref      # more stable than np.log(np.exp(lw - ref).sum(0)) + ref

    ## return
    return lw - lsum
//  [python]
def normalize_log_weights(lw):
    """Normalises a matrix of log-weights, row-wise"""

    ref = lw.max(0)
    lsum = sp.logsumexp(lw - ref, 0) + ref      # more stable than np.log(np.exp(lw - ref).sum(0)) + ref

    ## return
    return lw - lsum
//  [python]
def compute_rel_abun(x, denom):
    if x.ndim == 1:
        x = np.expand_dims(x, axis=0)
    z = np.hstack((x, np.zeros((x.shape[0], 1))))
    p = np.exp(z - np.max(z, axis=1, keepdims=True))
    p /= np.sum(p, axis=1, keepdims=True)
    for i in range(p.shape[1]-1, denom, -1):
        tmp = np.copy(p[:,i-1])
        p[:,i-1] = np.copy(p[:,i])
        p[:,i] = tmp
    return p
//  [python]
def compute_rel_abun(x, denom):
    if x.ndim == 1:
        x = np.expand_dims(x, axis=0)
    z = np.hstack((x, np.zeros((x.shape[0], 1))))
    p = np.exp(z - logsumexp(z, axis=1, keepdims=True))
    p /= p.sum(axis=1, keepdims=True)
    for i in range(p.shape[1]-1, denom, -1):
        tmp = np.copy(p[:,i-1])
        p[:,i-1] = np.copy(p[:,i])
        p[:,i] = tmp
    return p
//  [python]
def generalized_pg_mixture_2nd(k, alphas, betas):
    '''Function used to compute 
    generalization 2 (eq. 47 of 1902.08831 ). 

    second way to calculate generalized pg mixture, 
    based on iterative sum
    '''

    iters=[np.array(i) for i in bars_and_stars_iterator(int(k), len(betas))]
    
    log_res=[]
    for it in iters:
        
        log_res.append(calc_pg(it, alphas, betas).sum())
    
    return s.logsumexp(log_res)
//  [python]
def generalized_pg_mixture_2nd(k, alphas, betas):
    '''Function used to compute 
    generalization 2 (eq. 47 of 1902.08831 ). 

    second way to calculate generalized pg mixture, 
    based on iterative sum
    '''

    iters=[np.array(i) for i in bars_and_stars_iterator(int(k), len(betas))]
    
    log_res=[]
    for it in iters:
        
        log_res.append(calc_pg(it, alphas, betas).sum())
    
    return scipy.special.logsumexp(log_res)
//  [python]
def log_likelihood(self,x, K_extra=1):
        """
        Estimate the log likelihood with samples from
         the model. Draw k_extra components which were not populated by
         the current model in order to create a truncated approximate
         mixture model.
        """
        x = np.asarray(x)
        ks = self._get_occupied()
        K = len(ks)
        K_total = K + K_extra

        # Sample observation distributions given current labels
        obs_distns = []
        for k in range(K):
            o = copy.deepcopy(self.obs_distn)
            o.resample(data=self._get_data_withlabel(k))
            obs_distns.append(o)

        # Sample extra observation distributions from prior
        for k in range(K_extra):
            o = copy.deepcopy(self.obs_distn)
            o.resample()
            obs_distns.append(o)

        # Sample a set of weights
        weights = Categorical(alpha_0=self.alpha_0,
                              K=K_total,
                              weights=None)

        assert len(self.labels_list) == 1
        weights.resample(data=self.labels_list[0].z)

        # Now compute the log likelihood
        vals = np.empty((x.shape[0],K_total))
        for k in range(K_total):
            vals[:,k] = obs_distns[k].log_likelihood(x)

        vals += weights.log_likelihood(np.arange(K_total))
        assert not np.isnan(vals).any()
        return np.sum(np.log(np.sum(np.exp(vals),axis=1)))
//  [python]
def log_likelihood(self,x, K_extra=1):
        """
        Estimate the log likelihood with samples from
         the model. Draw k_extra components which were not populated by
         the current model in order to create a truncated approximate
         mixture model.
        """
        x = np.asarray(x)
        ks = self._get_occupied()
        K = len(ks)
        K_total = K + K_extra

        # Sample observation distributions given current labels
        obs_distns = []
        for k in range(K):
            o = copy.deepcopy(self.obs_distn)
            o.resample(data=self._get_data_withlabel(k))
            obs_distns.append(o)

        # Sample extra observation distributions from prior
        for k in range(K_extra):
            o = copy.deepcopy(self.obs_distn)
            o.resample()
            obs_distns.append(o)

        # Sample a set of weights
        weights = Categorical(alpha_0=self.alpha_0,
                              K=K_total,
                              weights=None)

        assert len(self.labels_list) == 1
        weights.resample(data=self.labels_list[0].z)

        # Now compute the log likelihood
        vals = np.empty((x.shape[0],K_total))
        for k in range(K_total):
            vals[:,k] = obs_distns[k].log_likelihood(x)

        vals += weights.log_likelihood(np.arange(K_total))
        assert not np.isnan(vals).any()
        return logsumexp(vals,axis=1).sum()
//  [python]
def evidence(ns_run, logw=None, simulate=False):
    r"""Bayesian evidence :math:`\log \mathcal{Z}`.

    Parameters
    ----------
    ns_run: dict
        Nested sampling run dict (see the data_processing module
        docstring for more details).
    logw: None or 1d numpy array, optional
        Log weights of samples.
    simulate: bool, optional
        Passed to ns_run_utils.get_logw if logw needs to be
        calculated.

    Returns
    -------
    float
    """
    if logw is None:
        logw = nestcheck.ns_run_utils.get_logw(ns_run, simulate=simulate)
    return np.exp(scipy.misc.logsumexp(logw))
//  [python]
def evidence(ns_run, logw=None, simulate=False):
    r"""Bayesian evidence :math:`\log \mathcal{Z}`.

    Parameters
    ----------
    ns_run: dict
        Nested sampling run dict (see the data_processing module
        docstring for more details).
    logw: None or 1d numpy array, optional
        Log weights of samples.
    simulate: bool, optional
        Passed to ns_run_utils.get_logw if logw needs to be
        calculated.

    Returns
    -------
    float
    """
    if logw is None:
        logw = nestcheck.ns_run_utils.get_logw(ns_run, simulate=simulate)
    return np.exp(scipy.special.logsumexp(logw))
//  [python]
def logz(ns_run, logw=None, simulate=False):
    r"""Natural log of Bayesian evidence :math:`\log \mathcal{Z}`.

    Parameters
    ----------
    ns_run: dict
        Nested sampling run dict (see the data_processing module
        docstring for more details).
    logw: None or 1d numpy array, optional
        Log weights of samples.
    simulate: bool, optional
        Passed to ns_run_utils.get_logw if logw needs to be
        calculated.

    Returns
    -------
    float
    """
    if logw is None:
        logw = nestcheck.ns_run_utils.get_logw(ns_run, simulate=simulate)
    return s.logsumexp(logw)
//  [python]
def logz(ns_run, logw=None, simulate=False):
    r"""Natural log of Bayesian evidence :math:`\log \mathcal{Z}`.

    Parameters
    ----------
    ns_run: dict
        Nested sampling run dict (see the data_processing module
        docstring for more details).
    logw: None or 1d numpy array, optional
        Log weights of samples.
    simulate: bool, optional
        Passed to ns_run_utils.get_logw if logw needs to be
        calculated.

    Returns
    -------
    float
    """
    if logw is None:
        logw = nestcheck.ns_run_utils.get_logw(ns_run, simulate=simulate)
    return scipy.special.logsumexp(logw)
//  [python]
def log_marginal_likelihood(self, *, x, all_g, z_g=None):
        """Computes the evidence log p(x) by exhaustive enumeration

        Args:
            x: observations [N, d]
            all_g: list of all possible DAGs with d nodes (igraph.Graph objects)
            
        Returns:
            [1,]
        """

        # log p(x, G)
        log_prob_obs_g = np.zeros(len(all_g))

        # normalizing constant for log p(G) using exhaustive normalization
        if z_g is None:
            z_g = self.g_dist.log_normalization_constant(all_g=all_g)

        # log p(x, G)
        for i, g in enumerate(tqdm.tqdm(all_g, desc='p(X) log_marginal_likelihood', disable=not self.verbose)):

            # log p(x, G) = log (p(G)/Z) + log p(x | G)
            log_prob_obs_g[i] = self.g_dist.unnormalized_log_prob(g=g) - z_g \
                + self.log_marginal_likelihood_given_g(g=g, x=x)

        # log p(x) = log(sum_G exp(log p(x, G)))
        return s.logsumexp(log_prob_obs_g)
//  [python]
def log_marginal_likelihood(self, *, x, all_g, z_g=None):
        """Computes the evidence log p(x) by exhaustive enumeration

        Args:
            x: observations [N, d]
            all_g: list of all possible DAGs with d nodes (igraph.Graph objects)
            
        Returns:
            [1,]
        """

        # log p(x, G)
        log_prob_obs_g = np.zeros(len(all_g))

        # normalizing constant for log p(G) using exhaustive normalization
        if z_g is None:
            z_g = self.g_dist.log_normalization_constant(all_g=all_g)

        # log p(x, G)
        for i, g in enumerate(tqdm.tqdm(all_g, desc='p(X) log_marginal_likelihood', disable=not self.verbose)):

            # log p(x, G) = log (p(G)/Z) + log p(x | G)
            log_prob_obs_g[i] = self.g_dist.unnormalized_log_prob(g=g) - z_g \
                + self.log_marginal_likelihood_given_g(g=g, x=x)

        # log p(x) = log(sum_G exp(log p(x, G)))
        return scipy.special.logsumexp(log_prob_obs_g)
//  [python]
def logistic_mixture_logpdf1(params, datum):
    # params are assumed to be normalized
    structured_params = params.reshape((-1, 3))
    component_scores = []
    for (loc, scale, component_prob) in structured_params:
        component_scores.append(
            logistic_logpdf(datum, loc, scale) + np.log(component_prob)
        )
    return scipy.misc.logsumexp(np.array(component_scores))
//  [python]
def logistic_mixture_logpdf1(params, datum):
    # params are assumed to be normalized
    structured_params = params.reshape((-1, 3))
    component_scores = []
    for (loc, scale, component_prob) in structured_params:
        component_scores.append(
            logistic_logpdf(datum, loc, scale) + np.log(component_prob)
        )
    return scipy.special.logsumexp(np.array(component_scores))
//  [python]
def _calc_val(self, trajectories):
        costs = trajectories["costs"].copy()
        actions = trajectories["actions"].copy()
        delta = actions - self.mean_action[None, :, :]
        
        traj_costs = cost_to_go(costs, self.gamma_seq)[:,0]
        control_costs = self._control_costs(delta)
        total_costs = traj_costs.copy() + self.lam * control_costs.copy()
        
        c = (-1.0/self.lam) * total_costs.copy()
        cmax = np.max(c)
        c -= cmax
        c = np.exp(c)
        val = cmax + np.log(np.sum(c)) - np.log(c.shape[0])
        val = -self.lam * val

        return val
//  [python]
def _calc_val(self, trajectories):
        costs = trajectories["costs"].copy()
        actions = trajectories["actions"].copy()
        delta = actions - self.mean_action[None, :, :]
        
        traj_costs = cost_to_go(costs, self.gamma_seq)[:,0]
        control_costs = self._control_costs(delta)
        total_costs = traj_costs.copy() + self.lam * control_costs.copy()
        
		# calculate log-sum-exp
        # c = (-1.0/self.lam) * total_costs.copy()
        # cmax = np.max(c)
        # c -= cmax
        # c = np.exp(c)
        # val1 = cmax + np.log(np.sum(c)) - np.log(c.shape[0])
        # val1 = -self.lam * val1

        val = -self.lam * scipy.special.logsumexp((-1.0/self.lam) * total_costs, b=(1.0/total_costs.shape[0]))
        return val
//  [python]
def _retraction(self, X, G):
        X = np.maximum(X, self.epsilon)
        norm = np.log(np.sum(np.exp((G - X) / X), axis=0))
        Y = np.log(X) + (G / X)
        if self.strict_simplex:
            Y -= norm[np.newaxis, :]
        else:
            mask = (norm > 0)
            Y[:, mask] -= norm[np.newaxis, mask]
        Y = np.exp(Y)

        mask = np.any(np.isnan(Y), axis=0)
        Y[:, mask] = X[:, mask]

        Y = np.maximum(Y, self.epsilon)

        if np.any(np.isnan(Y)):
            warnings.warn('Invalid value encountered during retraction mapping onto multinomial manifold')
        return np.nan_to_num(Y)
//  [python]
def _retraction(self, X, G):
        X = np.maximum(X, self.epsilon)
        norm = scipy.special.logsumexp(G / X, b=X, axis=0)
        Y = np.log(X) + (G / X)
        if self.strict_simplex:
            Y -= norm[np.newaxis, :]
        else:
            mask = (norm > 0)
            Y[:, mask] -= norm[np.newaxis, mask]
        Y = np.exp(Y)

        mask = np.any(np.isnan(Y), axis=0)
        Y[:, mask] = X[:, mask]

        Y = np.maximum(Y, self.epsilon)

        if np.any(np.isnan(Y)):
            warnings.warn('Invalid value encountered during retraction mapping onto multinomial manifold')
        return np.nan_to_num(Y)
//  [python]
def logmeanexp(x: Union[to.Tensor, np.ndarray], dim: int = 0) -> Union[to.Tensor, np.ndarray]:
    r"""
    Numerically stable way to compute $\log \left( \frac{1}{N} \sum_{i=1}^N \exp(x) \right)$

    :param x: input tensor
    :param dim: dimension to compute the logmeanexp` along
    :return: natural logarithm of the mean of the exponentiated values
    """
    if isinstance(x, to.Tensor):
        return to.logsumexp(x, dim=dim) - to.log(to.tensor(x.shape[dim], dtype=to.get_default_dtype()))
    elif isinstance(x, np.ndarray):
        return s.logsumexp(x, axis=dim) - np.log(np.array(x.shape[dim], dtype=np.float32))  # warning can be ignored
    else:
        raise pyrado.TypeErr(given=x, expected_type=[to.Tensor, np.ndarray])
//  [python]
def logmeanexp(x: Union[to.Tensor, np.ndarray], dim: int = 0) -> Union[to.Tensor, np.ndarray]:
    r"""
    Numerically stable way to compute $\log \left( \frac{1}{N} \sum_{i=1}^N \exp(x) \right)$

    :param x: input tensor
    :param dim: dimension to compute the logmeanexp` along
    :return: natural logarithm of the mean of the exponentiated values
    """
    if isinstance(x, to.Tensor):
        return to.logsumexp(x, dim=dim) - to.log(to.tensor(x.shape[dim], dtype=to.get_default_dtype()))
    elif isinstance(x, np.ndarray):
        return logsumexp(x, axis=dim) - np.log(np.array(x.shape[dim], dtype=np.float32))  # warning can be ignored
    else:
        raise pyrado.TypeErr(given=x, expected_type=[to.Tensor, np.ndarray])
//  [python]
def predict_proba(self, X):
        """
        Compute calibrated posterior probabilities for a given array of posterior probabilities from an arbitrary
        classifier.
        Parameters
        ----------
        X : array-like, shape (n_samples, n_classes)
            The uncalibrated posterior probabilities.
        Returns
        -------
        P : array, shape (n_samples, n_classes)
            The predicted probabilities.
        """
        if X.ndim == 1:
            raise ValueError("Calibration data must have shape (n_samples, n_classes).")
        elif np.shape(X)[1] == 2:
            check_is_fitted(self, ["binnings", "log_scores", "prob_class_1", "T"])

            # Find bin for all binnings and the associated empirical accuracy
            posterior_prob_binnings = np.zeros(shape=[np.shape(X)[0], len(self.binnings)])
            for i, binning in enumerate(self.binnings):
                bin_ids = np.searchsorted(binning, X[:, 1])
                bin_ids = np.clip(bin_ids, a_min=0, a_max=len(binning) - 1)  # necessary if X is out of range
                posterior_prob_binnings[:, i] = [self.prob_class_1[i][j] for j in (bin_ids - 1)]

            # Computed score-weighted average
            norm_weights = np.exp(np.array(self.log_scores) - np.max(self.log_scores) - np.log(np.sum(np.exp(np.array(self.log_scores) - np.max(self.log_scores)))))
            posterior_prob = np.sum(posterior_prob_binnings * norm_weights, axis=1)

            # Compute probability for other class
            return np.column_stack([1 - posterior_prob, posterior_prob])
        elif np.shape(X)[1] > 2:
            check_is_fitted(self, "onevsrest_calibrator_")
            return self.onevsrest_calibrator_.predict_proba(X)
//  [python]
def predict_proba(self, X):
        """
        Compute calibrated posterior probabilities for a given array of posterior probabilities from an arbitrary
        classifier.
        Parameters
        ----------
        X : array-like, shape (n_samples, n_classes)
            The uncalibrated posterior probabilities.
        Returns
        -------
        P : array, shape (n_samples, n_classes)
            The predicted probabilities.
        """
        if X.ndim == 1:
            raise ValueError("Calibration data must have shape (n_samples, n_classes).")
        elif np.shape(X)[1] == 2:
            check_is_fitted(self, ["binnings", "log_scores", "prob_class_1", "T"])

            # Find bin for all binnings and the associated empirical accuracy
            posterior_prob_binnings = np.zeros(shape=[np.shape(X)[0], len(self.binnings)])
            for i, binning in enumerate(self.binnings):
                bin_ids = np.searchsorted(binning, X[:, 1])
                bin_ids = np.clip(bin_ids, a_min=0, a_max=len(binning) - 1)  # necessary if X is out of range
                posterior_prob_binnings[:, i] = [self.prob_class_1[i][j] for j in (bin_ids - 1)]

            # Computed score-weighted average
            norm_weights = np.exp(np.array(self.log_scores) - scipy.special.logsumexp(self.log_scores))
            posterior_prob = np.sum(posterior_prob_binnings * norm_weights, axis=1)

            # Compute probability for other class
            return np.column_stack([1 - posterior_prob, posterior_prob])
        elif np.shape(X)[1] > 2:
            check_is_fitted(self, "onevsrest_calibrator_")
            return self.onevsrest_calibrator_.predict_proba(X)
//  [python]
def get_mc_entropy(mx, var):
    n, d     = mx.shape
    mx2 = mx + rnd.standard_normal(mx.shape, dtype=np.float32)*np.sqrt(var)

    dist_norm = pairwise_distance2_np(mx, mx2) / (-2.0 * var)  
    
    const     =  - 0.5  * d * np.log(2.0 * np.pi * var) - np.log(n)
    logprobs  = const + s.logsumexp(dist_norm , axis=0) 

    return -np.mean(logprobs)
//  [python]
def get_mc_entropy(mx, var):
    n, d     = mx.shape
    mx2 = mx + rnd.standard_normal(mx.shape, dtype=np.float32)*np.sqrt(var)

    dist_norm = pairwise_distance2_np(mx, mx2) / (-2.0 * var)  
    
    const     =  - 0.5  * d * np.log(2.0 * np.pi * var) - np.log(n)
    logprobs  = const + scipy.special.logsumexp(dist_norm , axis=0) 

    return -np.mean(logprobs)
//  [python]
def eval(self, x, ii=None, log=True, debug=False):
        """
        Evaluates the mog pdf.
        x: rows are inputs to evaluate at
        ii: a list of indices specifying which marginal to evaluate;
                   if None, the joint pdf is evaluated
        log: if True, the log pdf is evaluated
        :return: pdf or log pdf
        """
        ps = np.array([self.xs[ix].eval(x, ii, log)
                       for ix in range(len(self.a))]).T
        if log:
            res = np.log(np.sum(np.exp(ps + np.log(self.a)), axis=1))
        else:
            res = np.dot(ps, self.a)
        if debug:
            print('weights\n', self.a, '\nps\n', ps, '\nres\n', res)
        return res
//  [python]
def eval(self, x, ii=None, log=True, debug=False):
        """
        Evaluates the mog pdf.
        x: rows are inputs to evaluate at
        ii: a list of indices specifying which marginal to evaluate;
                   if None, the joint pdf is evaluated
        log: if True, the log pdf is evaluated
        :return: pdf or log pdf
        """
        ps = np.array([self.xs[ix].eval(x, ii, log)
                       for ix in range(len(self.a))]).T
        if log:
            res = scipy.special.logsumexp(ps + np.log(self.a), axis=1)
        else:
            res = np.dot(ps, self.a)
        if debug:
            print('weights\n', self.a, '\nps\n', ps, '\nres\n', res)
        return res
//  [python]
def __mul__(self, other):
        """Multiplies by a single Gaussian."""
        assert isinstance(other, Gaussian)
        ys = [x * other for x in self.xs]
        lcs = np.empty_like(self.a)
        for i, (x, y) in enumerate(zip(self.xs, ys)):
            lcs[i] = x.logdetP + other.logdetP - y.logdetP
            lcs[i] -= np.dot(x.m, np.dot(x.P, x.m))
            lcs[i] += np.dot(other.m, np.dot(other.P, other.m))
            lcs[i] -= np.dot(y.m, np.dot(y.P, y.m))
            lcs[i] *= 0.5
        la = np.log(self.a) + lcs
        la -= logsumexp(la)
        a = np.exp(la)
        return MoG(a=a, xs=ys)
//  [python]
def __div__(self, other):
        """Divides by a single Gaussian."""
        assert isinstance(other, Gaussian)
        ys = [x / other for x in self.xs]
        lcs = np.empty_like(self.a)
        for i, (x, y) in enumerate(zip(self.xs, ys)):
            lcs[i] = x.logdetP - other.logdetP - y.logdetP
            lcs[i] -= np.dot(x.m, np.dot(x.P, x.m))
            lcs[i] -= np.dot(other.m, np.dot(other.P, other.m))
            lcs[i] -= np.dot(y.m, np.dot(y.P, y.m))
            lcs[i] *= 0.5
        la = np.log(self.a) + lcs
        la -= logsumexp(la)
        a = np.exp(la)
        return MoG(a=a, xs=ys)
//  [python]
def expected_responsibilities(self, obs):
        log_prob = self.expected_log_complete_likelihood(obs)
        resp = np.exp(log_prob - logsumexp(log_prob, axis=0, keepdims=True))
        return resp
//  [python]
def get_log_likelihood(fx, normalizer=None, indifference_threshold = 0.0):
    """
    fx: (num_sample,num_data,num_choice)
    """
    num_sample = fx.shape[0]
    num_choice = fx.shape[2]
    num_data = fx.shape[1]

    if normalizer is None:
        normalizer = num_sample

    mask = np.eye(num_choice)
    indiff_mat = (1.0 - mask) * indifference_threshold
    
    fx = np.expand_dims(fx, axis=-1)
    # (num_sample, num_data, num_choice, 1)

    selected_fx = np.sum(fx * mask, axis=-2)
    base_fx = fx + indiff_mat

    log_likelihood_preference = selected_fx - s.logsumexp(base_fx, axis=-2)
    # (num_sample, num_data, num_choice)

    prob_indifference = 1.0 - np.sum(np.exp(log_likelihood_preference), axis=-1, keepdims=True)
    prob_indifference = np.clip(prob_indifference, a_min=1e-100, a_max=1.0 - 1e-100)
    log_likelihood_indifference = np.log(prob_indifference)
    # (num_sample, num_data, 1)

    log_likelihood = np.concatenate([log_likelihood_preference, log_likelihood_indifference], axis=-1)
    # (num_sample, num_data, num_choice + 1)
    
    log_likelihood = s.logsumexp( log_likelihood - np.log(normalizer), axis=0 )
    # (num_data, num_choice + 1)

    return log_likelihood.T
//  [python]
def get_log_likelihood(fx, normalizer=None, indifference_threshold = 0.0):
    """
    fx: (num_sample,num_data,num_choice)
    """
    num_sample = fx.shape[0]
    num_choice = fx.shape[2]
    num_data = fx.shape[1]

    if normalizer is None:
        normalizer = num_sample

    mask = np.eye(num_choice)
    indiff_mat = (1.0 - mask) * indifference_threshold
    
    fx = np.expand_dims(fx, axis=-1)
    # (num_sample, num_data, num_choice, 1)

    selected_fx = np.sum(fx * mask, axis=-2)
    base_fx = fx + indiff_mat

    log_likelihood_preference = selected_fx - spmc.logsumexp(base_fx, axis=-2)
    # (num_sample, num_data, num_choice)

    prob_indifference = 1.0 - np.sum(np.exp(log_likelihood_preference), axis=-1, keepdims=True)
    prob_indifference = np.clip(prob_indifference, a_min=1e-100, a_max=1.0 - 1e-100)
    log_likelihood_indifference = np.log(prob_indifference)
    # (num_sample, num_data, 1)

    log_likelihood = np.concatenate([log_likelihood_preference, log_likelihood_indifference], axis=-1)
    # (num_sample, num_data, num_choice + 1)
    
    log_likelihood = spmc.logsumexp( log_likelihood - np.log(normalizer), axis=0 )
    # (num_data, num_choice + 1)

    return log_likelihood.T
//  [python]
def logpdf(self, x):
        scores = []
        for (c, p) in zip(self.components, self.probs):
            scores.append(c.logpdf(x) + np.log(p))
        return scipy.special.logsumexp(np.array(scores))
//  [python]
def sampling_rule(self, temperature=None):
        """
            Select an action from the node.
            - if exploration is wanted with some temperature, follow the selection strategy.
            - else, select the action with maximum visit count

        :param temperature: the exploration parameter, positive or zero
        :return: the selected action
        """

        if self.children:
            all_childrens = list(self.children.values())
            a_dim = len(all_childrens)
            if (self.count < a_dim):
                return self.count

            tem = (temperature * a_dim) / np.log(1+self.count)

            average_pro = np.ones(a_dim) / a_dim
            
            if tem < 1:
                all_q_values = np.array([child.value for child in all_childrens])
                expqsum = np.log(np.sum(np.exp(all_q_values)))
                log_pro = all_q_values - expqsum
                pro = np.exp(log_pro)
                pro = (1 - tem) * pro + tem * average_pro
            else:
                pro = average_pro
            pro = np.cumsum(pro)
            p_sample = 0.999999*np.random.rand(1)
            idx = np.searchsorted(pro, p_sample)
            return idx[0]
        else:
            return None
//  [python]
def sampling_rule(self, temperature=None):
        """
            Select an action from the node.
            - if exploration is wanted with some temperature, follow the selection strategy.
            - else, select the action with maximum visit count

        :param temperature: the exploration parameter, positive or zero
        :return: the selected action
        """

        if self.children:
            all_childrens = list(self.children.values())
            a_dim = len(all_childrens)
            if (self.count < a_dim):
                return self.count

            tem = (temperature * a_dim) / np.log(1+self.count)

            average_pro = np.ones(a_dim) / a_dim
            
            if tem < 1:
                all_q_values = np.array([child.value for child in all_childrens])
                expqsum = scipy.special.logsumexp(all_q_values)
                log_pro = all_q_values - expqsum
                pro = np.exp(log_pro)
                pro = (1 - tem) * pro + tem * average_pro
            else:
                pro = average_pro
            pro = np.cumsum(pro)
            p_sample = 0.999999*np.random.rand(1)
            idx = np.searchsorted(pro, p_sample)
            return idx[0]
        else:
            return None
//  [python]
def update_branch(self, total_reward):
        """
            Update the whole branch from this node to the root with the total reward of the corresponding trajectory.

        :param total_reward: the total reward obtained through a trajectory passing by this node
        """
        self.update(total_reward)
        if self.parent:
            all_childrens = list(self.parent.children.values())
            all_q_values = np.array([child.value for child in all_childrens])
            all_pre_values = np.array([np.log(child.prior) for child in all_childrens])
            logexpqsum = np.log(np.sum(np.exp(all_q_values + all_pre_values)))

            self.parent.update_branch(logexpqsum)
        elif (self.planner.config["Plot"]==1):
            all_childrens = list(self.children.values())
            all_q_values = np.array([child.value for child in all_childrens])
            all_pre_values = np.array([np.log(child.prior) for child in all_childrens])
            logexpqsum = np.log(np.sum(np.exp(all_q_values + all_pre_values)))
            log_pro = all_q_values + all_pre_values - logexpqsum
            pro = np.exp(log_pro)