jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def plot_eta_stats(eta_stats, alpha, threshold=0.01):
  K_vals = list(sorted(set([K for (K, S) in eta_stats.keys()])))
  S_vals = list(sorted(set([S for (K, S) in eta_stats.keys()])))
  K_map = {K: kidx for kidx, K in enumerate(K_vals)}
  S_map = {S: sidx for sidx, S in enumerate(S_vals)}

  below_thresh_empirical   = np.full((len(K_vals), len(S_vals)), np.nan)
  below_thresh_lol = {}
  below_thresh_theoretical = np.copy(below_thresh_empirical)

  for (K, S), biggest_eta in eta_stats.items():
    num_trees = len(biggest_eta)
    assert biggest_eta.shape == (num_trees, K)
    bt = np.sum(biggest_eta < threshold, axis=1)
    frac = bt/K
    assert len(frac) == num_trees
    below_thresh_empirical[K_map[K],S_map[S]] = np.mean(frac)
    below_thresh_lol[(K,S)] = frac

  for K in K_vals:
    for S in S_vals:
      prob = s.betainc(alpha, (K-1)*alpha, threshold)**S
      below_thresh_theoretical[K_map[K],S_map[S]] = prob

  fig_emp_hm = make_heatmap(
    S_vals,
    K_vals,
    below_thresh_empirical,
    'Cancer samples',
    'Populations',
    'Mean proportion of populations with η < 0.01<br>across all cancer samples (empirical)',
  )
  fig_thr_hm = make_heatmap(
    S_vals,
    K_vals,
    below_thresh_theoretical,
    'Cancer samples',
    'Populations',
    'Probability that η < 0.01<br>across all cancer samples (theoretical)',
  )

  fig_emp_sp = make_subplots(
    rows = len(K_vals),
    cols = 1,
    shared_xaxes = True,
    vertical_spacing = 0.05,
    row_titles = ['%s subclones' % K for K in K_vals],
    x_title = 'Cancer samples',
    y_title = 'Proportion of populations with η < 0.01<br>across all cancer samples',
  )
  for K in K_vals:
    for S in S_vals:
      if (K, S) not in below_thresh_lol:
        continue
      fig_emp_sp.add_trace(go.Box(y=below_thresh_lol[(K,S)], name=S, boxmean=True), row=K_map[K]+1, col=1)
    fig_emp_sp.add_trace(
      go.Scatter(
        x = S_vals,
        y = below_thresh_theoretical[K_map[K]],
        mode = 'markers',
        marker = {'size': 22, 'opacity': 0.65, 'color': '#000000'},
      ),
      row=K_map[K]+1,
      col=1
    )
  fig_emp_sp.update_xaxes(type = 'category')
  fig_emp_sp.update_yaxes(tickformat = ',.0%', range = (0,1))
  fig_emp_sp.update_layout(showlegend = False)

  fig_thr_sp = make_subplots(
    rows = len(K_vals),
    cols = 1,
    shared_xaxes = True,
    vertical_spacing = 0.09,
    x_title = 'Cancer samples',
    y_title = 'Proportion of populations with η < 0.01 (theoretical)',
  )
  for K in K_vals:
    fig_thr_sp.add_trace(
      go.Bar(x=S_vals, y=below_thresh_theoretical[K_map[K]]),
      row = len(K_vals) - K_map[K],
      col = 1,
    )
  fig_thr_sp.update_xaxes(type = 'category')
  fig_thr_sp.update_yaxes(tickformat = ',.0%', range = (0,1))

  return fig_emp_hm + fig_thr_hm + to_html(fig_emp_sp) + to_html(fig_thr_sp)
//  [python]
def plot_eta_stats(eta_stats, alpha, threshold=0.01):
  K_vals = list(sorted(set([K for (K, S) in eta_stats.keys()])))
  S_vals = list(sorted(set([S for (K, S) in eta_stats.keys()])))
  K_map = {K: kidx for kidx, K in enumerate(K_vals)}
  S_map = {S: sidx for sidx, S in enumerate(S_vals)}

  below_thresh_empirical   = np.full((len(K_vals), len(S_vals)), np.nan)
  below_thresh_lol = {}
  below_thresh_theoretical = np.copy(below_thresh_empirical)

  for (K, S), biggest_eta in eta_stats.items():
    num_trees = len(biggest_eta)
    assert biggest_eta.shape == (num_trees, K)
    bt = np.sum(biggest_eta < threshold, axis=1)
    frac = bt/K
    assert len(frac) == num_trees
    below_thresh_empirical[K_map[K],S_map[S]] = np.mean(frac)
    below_thresh_lol[(K,S)] = frac

  for K in K_vals:
    for S in S_vals:
      prob = scipy.special.betainc(alpha, (K-1)*alpha, threshold)**S
      below_thresh_theoretical[K_map[K],S_map[S]] = prob

  fig_emp_hm = make_heatmap(
    S_vals,
    K_vals,
    below_thresh_empirical,
    'Cancer samples',
    'Populations',
    'Mean proportion of populations with η < 0.01<br>across all cancer samples (empirical)',
  )
  fig_thr_hm = make_heatmap(
    S_vals,
    K_vals,
    below_thresh_theoretical,
    'Cancer samples',
    'Populations',
    'Probability that η < 0.01<br>across all cancer samples (theoretical)',
  )

  fig_emp_sp = make_subplots(
    rows = len(K_vals),
    cols = 1,
    shared_xaxes = True,
    vertical_spacing = 0.05,
    row_titles = ['%s subclones' % K for K in K_vals],
    x_title = 'Cancer samples',
    y_title = 'Proportion of populations with η < 0.01<br>across all cancer samples',
  )
  for K in K_vals:
    for S in S_vals:
      if (K, S) not in below_thresh_lol:
        continue
      fig_emp_sp.add_trace(go.Box(y=below_thresh_lol[(K,S)], name=S, boxmean=True), row=K_map[K]+1, col=1)
    fig_emp_sp.add_trace(
      go.Scatter(
        x = S_vals,
        y = below_thresh_theoretical[K_map[K]],
        mode = 'markers',
        marker = {'size': 22, 'opacity': 0.65, 'color': '#000000'},
      ),
      row=K_map[K]+1,
      col=1
    )
  fig_emp_sp.update_xaxes(type = 'category')
  fig_emp_sp.update_yaxes(tickformat = ',.0%', range = (0,1))
  fig_emp_sp.update_layout(showlegend = False)

  fig_thr_sp = make_subplots(
    rows = len(K_vals),
    cols = 1,
    shared_xaxes = True,
    vertical_spacing = 0.09,
    x_title = 'Cancer samples',
    y_title = 'Proportion of populations with η < 0.01 (theoretical)',
  )
  for K in K_vals:
    fig_thr_sp.add_trace(
      go.Bar(x=S_vals, y=below_thresh_theoretical[K_map[K]]),
      row = len(K_vals) - K_map[K],
      col = 1,
    )
  fig_thr_sp.update_xaxes(type = 'category')
  fig_thr_sp.update_yaxes(tickformat = ',.0%', range = (0,1))

  return fig_emp_hm + fig_thr_hm + to_html(fig_emp_sp) + to_html(fig_thr_sp)
//  [python]
def log_sf(self,x):
        scalar = not isinstance(x,np.ndarray)
        x = np.atleast_1d(x)
        errs = np.seterr(divide='ignore')
        ret = np.log(s.betainc(x+1,self.r,self.p))
        np.seterr(**errs)
        ret[x < 0] = np.log(1.)
        if scalar:
            return ret[0]
        else:
            return ret
//  [python]
def log_sf(self,x):
        scalar = not isinstance(x,np.ndarray)
        x = np.atleast_1d(x)
        errs = np.seterr(divide='ignore')
        ret = np.log(special.betainc(x+1,self.r,self.p))
        np.seterr(**errs)
        ret[x < 0] = np.log(1.)
        if scalar:
            return ret[0]
        else:
            return ret
//  [python]
def make_table(cls, n_interp=200, n_dims=50):
        """Save lookup table for ``_betaincinv22.lookup``."""
        import scipy.special  # pylint: disable=import-outside-toplevel

        rng = np.random.RandomState(0)

        n_dims_log = int(0.5 * n_dims)
        dims_lin = np.arange(1, n_dims - n_dims_log + 1)
        dims_log = np.round(
            np.logspace(np.log10(dims_lin[-1] + 1), 3, n_dims_log)
        ).astype(dims_lin.dtype)
        dims = np.unique(np.concatenate([dims_lin, dims_log]))

        x_table = []
        y_table = []
        for dim in dims:
            n_range = int(0.8 * n_interp)
            x0 = np.linspace(0, 1, n_interp - n_range)  # samples in the domain
            y0 = np.linspace(1e-16, 1 - 1e-7, n_range)  # samples in the range
            x1 = scipy.special.betainc(dim / 2.0, 0.5, y0)
            interp_x = np.unique(np.concatenate([x0, x1]))
            while len(interp_x) < n_interp:
                # add random points until we achieve the length
                interp_x = np.unique(
                    np.concatenate(
                        [interp_x, rng.uniform(size=n_interp - len(interp_x))]
                    )
                )

            interp_x.sort()
            assert interp_x.size == n_interp

            interp_y = scipy.special.betaincinv(dim / 2.0, 0.5, interp_x)
            x_table.append(interp_x)
            y_table.append(interp_y)

        x_table = np.asarray(x_table)
        y_table = np.asarray(y_table)

        np.savez(cls.path, dims=dims, x=x_table, y=y_table)
//  [python]
def cdf(pX, pR, pP):
    """
    Cumulative density function of a continuous generalization of NB distribution
    """
    # if pX == 0:
    # return 0
    return s.betainc(pR, pX + 1, pP)
//  [python]
def cdf(pX, pR, pP):
    """
    Cumulative density function of a continuous generalization of NB distribution
    """
    # if pX == 0:
    # return 0
    return special.betainc(pR, pX + 1, pP)
//  [python]
def _calc_logbf(ssm, omega1=1):
  V = ssm['var_reads']
  R = ssm['ref_reads']
  S = len(V)

  omega = [np.broadcast_to(omega1, S), ssm['omega_v']]
  logp = []

  for om in omega:
    integ = s.betainc(V+1, R+1, om)
    # Prevent division by zero.
    loginteg = np.log(np.maximum(1e-30, integ))
    logp.append(-np.log(om) + loginteg)

  logbf = logp[0] - logp[1]
  # Convert from nats to bits.
  logbf /= np.log(2)

  return logbf
//  [python]
def _calc_logbf(ssm, omega1=1):
  V = ssm['var_reads']
  R = ssm['ref_reads']
  S = len(V)

  omega = [np.broadcast_to(omega1, S), ssm['omega_v']]
  logp = []

  for om in omega:
    integ = scipy.special.betainc(V+1, R+1, om)
    # Prevent division by zero.
    loginteg = np.log(np.maximum(1e-30, integ))
    logp.append(-np.log(om) + loginteg)

  logbf = logp[0] - logp[1]
  # Convert from nats to bits.
  logbf /= np.log(2)

  return logbf
//  [python]
def cdf(self, x):
        return scipy.special.beta(self.alpha, self.beta, x)
//  [python]
def cdf(self, x):
        return scipy.special.betainc(self.alpha, self.beta, x)
//  [python]
def check_sample_mean(sm, v, n, popmean):
    # from stats.stats.ttest_1samp(a, popmean):
    # Calculates the t-obtained for the independent samples T-test on ONE group
    # of scores a, given a population mean.
    #
    # Returns: t-value, two-tailed prob
    df = n-1
    svar = ((n-1)*v) / float(df)    # looks redundant
    t = (sm-popmean) / np.sqrt(svar*(1.0/n))
    prob = s.betainc(0.5*df, 0.5, df/(df + t*t))

    # return t,prob
    npt.assert_(prob > 0.01, 'mean fail, t,prob = %f, %f, m, sm=%f,%f' %
                (t, prob, popmean, sm))
//  [python]
def check_sample_mean(sm, v, n, popmean):
    # from stats.stats.ttest_1samp(a, popmean):
    # Calculates the t-obtained for the independent samples T-test on ONE group
    # of scores a, given a population mean.
    #
    # Returns: t-value, two-tailed prob
    df = n-1
    svar = ((n-1)*v) / float(df)    # looks redundant
    t = (sm-popmean) / np.sqrt(svar*(1.0/n))
    prob = betainc(0.5*df, 0.5, df/(df + t*t))

    # return t,prob
    npt.assert_(prob > 0.01, 'mean fail, t,prob = %f, %f, m, sm=%f,%f' %
                (t, prob, popmean, sm))
//  [python]
def inverseCdf(self, x):
    """
      Inverse cumulative distribution function. The scipy implementation of this function does not match the behavior
      of the boost implementation, so we implement it manually here.

      @ In, x, float, point at which to evaluate the inverse cdf
      @ Out, inverseCdf, float, inverse cumulative distribution function at x
    """
    # Solve numerically using the complement of the incomplete beta function
    def func(a, y):
      """
        CDF defined with the complement of the incomplete beta function

        @ In, a, float, point at which to evaluate the CDF
        @ In, y, float, desired quantile value
        @ Out, diff, float, difference between the desired quantile value and the CDF evaluated at a
      """
      if a == self.dist.args[0]:
        q = 1
      else:
        q = 1 - scipy.special.betainc(a + 1, self.dist.args[0] - a, self.dist.args[1])
      diff = q - y
      return diff

    # The root finding problem being solved to get the inverse CDF is not vectorized, so we need to loop over the
    # elements of x if x is an array.
    if hasattr(x, '__len__'):
      # Uses the toms748 algorithm for root finding, as stated in the boost documentation at
      # https://www.boost.org/doc/libs/1_50_0/libs/math/doc/sf_and_dist/html/math_toolkit/dist/dist_ref/dists/binomial_dist.html
      roots = [scipy.optimize.root_scalar(func, bracket=[0, self.dist.args[0]], args=(xi), method='toms748').root for xi in x]
      # Round "out": down if less than 0.5, up if greater than 0.5
      iCdf = np.array([np.floor(root) if root < 0.5 else np.ceil(root) for root in roots])
    else:
      root = scipy.optimize.root_scalar(func, bracket=[0, self.dist.args[0]], args=(x), method='toms748').root
      iCdf = np.floor(root) if root < 0.5 else np.ceil(root)
    return iCdf
//  [python]
def get_correlation_and_pvalue(n_obs, sx, sxx, sxy, sy, syy):
    df = n_obs - 2
    d = np.sqrt(n_obs * sxx - sx * sx) * np.sqrt(n_obs * syy - sy * sy)[:, np.newaxis]
    r = (n_obs * sxy - sx * sy[:, np.newaxis]) / d
    r[d == 0] = 0
    r = r.clip(-1, 1)
    t_squared = r ** 2 * (df / ((1.0 - r) * (1.0 + r)))
    prob = special.betainc(
        0.5 * df, 0.5, np.fmin(np.asarray(df / (df + t_squared)), 1.0)
    )
    prob[abs(r) == 1] = 0
    return r, prob
//  [python]
def K(r, R, beta):
        """Equation A16 im Mamon & Lokas for constant anisotropy.

        :param r: 3d radius
        :param R: projected 2d radius
        :param beta: anisotropy, float >-0.5
        :return: K(r, R, beta)
        """
        u = r / R
        k = np.sqrt(1 - 1.0 / u**2) / (1.0 - 2 * beta) + np.sqrt(
            np.pi
        ) / 2 * special.gamma(beta - 1.0 / 2) / special.gamma(beta) * (
            3.0 / 2 - beta
        ) * u ** (
            2 * beta - 1.0
        ) * (
            1 - special.beta(beta + 1.0 / 2, 1.0 / 2) - special.betainc(1.0/u**2, beta + 1.0 / 2, 1.0 / 2)
        )
        return k
//  [python]
def K(r, R, beta):
        """Equation A16 im Mamon & Lokas for constant anisotropy.

        :param r: 3d radius
        :param R: projected 2d radius
        :param beta: anisotropy, float >-0.5
        :return: K(r, R, beta)
        """
        u = r / R
        k = np.sqrt(1 - 1.0 / u**2) / (1.0 - 2 * beta) + np.sqrt(
            np.pi
        ) / 2 * special.gamma(beta - 1.0 / 2) / special.gamma(beta) * (
            3.0 / 2 - beta
        ) * u ** (
            2 * beta - 1.0
        ) * (
            1 - special.betainc(beta + 1.0 / 2, 1.0 / 2, 1.0 / u**2)
        )
        return k
//  [python]
def test_betainc():
    a, b, x = at.scalars("a", "b", "x")
    res = betainc(a, b, x)
    test_func = function([a, b, x], res, mode=Mode("py"))
    assert np.isclose(test_func(15, 10, 0.7), s.betainc(15, 10, 0.7))
//  [python]
def test_betainc():
    a, b, x = at.scalars("a", "b", "x")
    res = betainc(a, b, x)
    test_func = function([a, b, x], res, mode=Mode("py"))
    assert np.isclose(test_func(15, 10, 0.7), sp.betainc(15, 10, 0.7))
//  [python]
def update_log_evidence(self):
        """Sum up all the model-specific evidences from the submodels"""
        self.log_evidence = scipy.misc.logsumexp([model.model_log_evidence
                                for model in self.model_universe])
//  [python]
def update_log_evidence(self):
        """Sum up all the model-specific evidences from the submodels"""
        self.log_evidence = scipy.special.logsumexp([model.model_log_evidence
                                for model in self.model_universe])
//  [python]
def incomplete_beta_integral(a, b, x):
    """Returns the incomplete beta function, namely:

    I_x(a,b) = 1/B(a,b)*(Integral(0,x) of t^(a-1)(1-t)^(b-1) dt),

    where a > 0, b > 0 and B(a,b) = G(a)*G(b)/(G(a+b)).
    G(a) denotes the gamma function of a.
    The standard broadcasting rules apply to a, b, and x.

    Parameters
    ----------
    a : array_like or float > 0
    b : array_like or float > 0
    x : array_like or float
        x will be made no greater than 1.0.

    Returns
    -------
    betai : ndarray
        The incomplete beta function for the parameters provided.
    """

    x = np.array(x, copy = True)
    x = np.where(x < 1.0, x, 1.0)  # values of x larger than 1.0 are clipped.

    return s.betainc(a, b, x)
//  [python]
def incomplete_beta_integral(a, b, x):
    """Returns the incomplete beta function, namely:

    I_x(a,b) = 1/B(a,b)*(Integral(0,x) of t^(a-1)(1-t)^(b-1) dt),

    where a > 0, b > 0 and B(a,b) = G(a)*G(b)/(G(a+b)).
    G(a) denotes the gamma function of a.
    The standard broadcasting rules apply to a, b, and x.

    Parameters
    ----------
    a : array_like or float > 0
    b : array_like or float > 0
    x : array_like or float
        x will be made no greater than 1.0.

    Returns
    -------
    betai : ndarray
        The incomplete beta function for the parameters provided.
    """

    x = np.array(x, copy = True)
    x = np.where(x < 1.0, x, 1.0)  # values of x larger than 1.0 are clipped.

    return betainc(a, b, x)
//  [python]
def save_negative_log_likelihood(self,t):
        """Get the negative log likelihood if you want to store it"""
        #DEBUG: Unclear why model and run length log distr has different length
        #       Maybe because of trimming?
        all_one_step_ahead_lklhoods_weighted = [
            (self.model_universe[m].one_step_ahead_predictive_log_probs + 
             self.model_and_run_length_log_distr[m,
                        self.model_and_run_length_log_distr[m,:]> -np.inf])
             #self.model_universe[m].retained_run_lengths])
             #np.where(self.model_and_run_length_log_distr[m,:]> -np.inf)])            
             for m in range(0, self.Q)]
        summed_up = -misc.logsumexp([item for entry in 
            all_one_step_ahead_lklhoods_weighted for item in entry])
        self.negative_log_likelihood.append(summed_up)
//  [python]
def save_negative_log_likelihood(self,t):
        """Get the negative log likelihood if you want to store it"""
        #DEBUG: Unclear why model and run length log distr has different length
        #       Maybe because of trimming?
        all_one_step_ahead_lklhoods_weighted = [
            (self.model_universe[m].one_step_ahead_predictive_log_probs + 
             self.model_and_run_length_log_distr[m,
                        self.model_and_run_length_log_distr[m,:]> -np.inf])
             #self.model_universe[m].retained_run_lengths])
             #np.where(self.model_and_run_length_log_distr[m,:]> -np.inf)])            
             for m in range(0, self.Q)]
        summed_up = -scipy.special.logsumexp([item for entry in 
            all_one_step_ahead_lklhoods_weighted for item in entry])
        self.negative_log_likelihood.append(summed_up)
//  [python]
def partial_eval(self, X, n_samples=5):
    if n_samples < 1000:
      res, iwae = self.sess.run(
          (self.lHat, self.iwae),
          feed_dict={self.x: X, self.n_samples: n_samples})
      res = [iwae] + res
    else:  # special case to handle OOM
      assert n_samples % 100 == 0, "When using large # of samples, it must be divisble by 100"
      res = []
      for i in xrange(int(n_samples/100)):
        logF, = self.sess.run(
            (self.logF,),
            feed_dict={self.x: X, self.n_samples: 100})
        res.append(logsumexp(logF, axis=1))
      res = [np.mean(logsumexp(res, axis=0) - np.log(n_samples))]
    return res
//  [python]
def partial_eval(self, X, n_samples=5):
    if n_samples < 1000:
      res, iwae = self.sess.run(
          (self.lHat, self.iwae),
          feed_dict={self.x: X, self.n_samples: n_samples})
      res = [iwae] + res
    else:  # special case to handle OOM
      assert n_samples % 100 == 0, "When using large # of samples, it must be divisble by 100"
      res = []
      for i in xrange(int(n_samples/100)):
        logF, = self.sess.run(
            (self.logF,),
            feed_dict={self.x: X, self.n_samples: 100})
        res.append(scipy.special.logsumexp(logF, axis=1))
      res = [np.mean(scipy.special.logsumexp(res, axis=0) - np.log(n_samples))]
    return res
//  [python]
def update_priors(self, t):
        """update the priors for each model in the model universe, provided 
        that the boolean auto_prior_update is true"""
        for (m, model) in zip(range(0,self.Q), self.model_universe):
            if model.auto_prior_update == True:
                """We need to weigh quantities acc. to rld"""
                model_specific_rld = (self.model_and_run_length_log_distr[m,:] -
                    misc.logsumexp(self.model_and_run_length_log_distr[m,:]))
                model.prior_update(t, model_specific_rld)
//  [python]
def update_priors(self, t):
        """update the priors for each model in the model universe, provided 
        that the boolean auto_prior_update is true"""
        for (m, model) in zip(range(0,self.Q), self.model_universe):
            if model.auto_prior_update == True:
                """We need to weigh quantities acc. to rld"""
                model_specific_rld = (self.model_and_run_length_log_distr[m,:] -
                    scipy.special.logsumexp(self.model_and_run_length_log_distr[m,:]))
                model.prior_update(t, model_specific_rld)
//  [python]
def save_negative_log_likelihood_fixed_pars(self, t):
        """For comparison with GP, fix sigma^2 at the MAP when doing the 
        one-step-ahead prediction."""
        all_one_step_ahead_lklhoods_weighted = [
            (self.model_universe[m].one_step_ahead_predictive_log_probs_fixed_pars + 
             self.model_and_run_length_log_distr[m,
                        self.model_and_run_length_log_distr[m,:]> -np.inf])
            for m in range(0, self.Q)]
             #self.model_universe[m].retained_run_lengths])
             #np.where(self.model_and_run_length_log_distr[m,:]> -np.inf)])            
             
        summed_up = -misc.logsumexp([item for entry in 
            all_one_step_ahead_lklhoods_weighted for item in entry])
        self.negative_log_likelihood_fixed_pars.append(summed_up)
//  [python]
def save_negative_log_likelihood_fixed_pars(self, t):
        """For comparison with GP, fix sigma^2 at the MAP when doing the 
        one-step-ahead prediction."""
        all_one_step_ahead_lklhoods_weighted = [
            (self.model_universe[m].one_step_ahead_predictive_log_probs_fixed_pars + 
             self.model_and_run_length_log_distr[m,
                        self.model_and_run_length_log_distr[m,:]> -np.inf])
            for m in range(0, self.Q)]
             #self.model_universe[m].retained_run_lengths])
             #np.where(self.model_and_run_length_log_distr[m,:]> -np.inf)])            
             
        summed_up = -scipy.special.logsumexp([item for entry in 
            all_one_step_ahead_lklhoods_weighted for item in entry])
        self.negative_log_likelihood_fixed_pars.append(summed_up)
//  [python]
def eval(self, data):
        llchildren = np.array([c.eval(data) for c in self.children])

        return scipy.misc.logsumexp(llchildren + self.log_weights, axis=0)
//  [python]
def eval(self, data):
        llchildren = np.array([c.eval(data) for c in self.children])

        return scipy.special.logsumexp(llchildren + self.log_weights, axis=0)
//  [python]
def sliding_correlation(a,b,W, method="df_corr", sig=10, plot=False):
    
    
    if method == "df_corr":
        
        df_coef = a.rolling(W).corr(b)
        deg_of_free = W - 2
        t_squared = df_coef**2 * (deg_of_free / ((1.0 - df_coef) * (1.0 + df_coef)))
        pval = special.betainc(0.5*deg_of_free, 0.5, deg_of_free/(deg_of_free + t_squared))
        
        if sig is not None:
            coef_sig = df_coef[pval < sig/100].max()
            
    else:
        # a,b are input arrays; W is window length
        time = a.index
        
        a = a.to_numpy(copy=False)
        b = b.to_numpy(copy=False)
        
        df_coef = pd.DataFrame(index=time, columns= ["coef"])
        am = uniform_filter(a.astype(float),W)
        bm = uniform_filter(b.astype(float),W)
    
        amc = am[W//2:-W//2+1]
        bmc = bm[W//2:-W//2+1]
    
        da = a[:,None]-amc
        db = b[:,None]-bmc
    
        # Get sliding mask of valid windows
        m,n = da.shape
        mask1 = np.arange(m)[:,None] >= np.arange(n)
        mask2 = np.arange(m)[:,None] < np.arange(n)+W
        mask = mask1 & mask2
        dam = (da*mask)
        dbm = (db*mask)
    
        ssAs = np.einsum('ij,ij->j',dam,dam)
        ssBs = np.einsum('ij,ij->j',dbm,dbm)
        D = np.einsum('ij,ij->j',dam,dbm)
        coeff = D/np.sqrt(ssAs*ssBs)
    
        n = W
        ab = n/2 - 1
        pval = 2*special.btdtr(ab, ab, 0.5*(1 - abs(np.float64(coeff))))
        df_coef.iloc[W-1:, 0] = coeff
        
        if sig is not None:
            coef_sig = np.min(coeff[pval < sig/100])
            
    
    if plot:
        fig, ax = plt.subplots(figsize=(15, 4))
        ax.set_ylabel('Correlation Coefficient', fontweight="bold", fontsize=20)
        ax.plot(df_coef, color="black")
        ax.xaxis.set_major_locator(YearLocator(5))
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
        ax.axhline(y=np.abs(coef_sig), linestyle="--", color="red", linewidth=2)
        ax.axhline(y=0, linestyle="--", color="grey", linewidth=2)
        ax.axhline(y=np.abs(coef_sig)*-1, linestyle="--", color="red", linewidth=2)
        plt.show()
            
    if sig is not None:
        return df_coef, np.abs(coef_sig)
    
    else:
        return df_coef
//  [python]
def logsoftmax(x):
    """
    Calculates the log softmax of x, or of the rows of x.
    :param x: vector or matrix
    :return: log softmax
    """

    x = np.asarray(x)

    if x.ndim == 1:
        return x - scipy.misc.logsumexp(x)

    elif x.ndim == 2:
        return x - scipy.misc.logsumexp(x, axis=1)[:, np.newaxis]

    else:
        raise ValueError('input must be either vector or matrix')
//  [python]
def logsoftmax(x):
    """
    Calculates the log softmax of x, or of the rows of x.
    :param x: vector or matrix
    :return: log softmax
    """

    x = np.asarray(x)

    if x.ndim == 1:
        return x - scipy.special.logsumexp(x)

    elif x.ndim == 2:
        return x - scipy.special.logsumexp(x, axis=1)[:, np.newaxis]

    else:
        raise ValueError('input must be either vector or matrix')
//  [python]
def importance_sample(target, proposal, n_samples, rng=np.random):
    """
    Importance sampling.
    :param target: target distribution
    :param proposal: proposal distribution
    :param n_samples: number of samples
    :param rng: random generator to use
    :return: samples, normalized log weights
    """

    xs = proposal.gen(n_samples, rng=rng)
    log_ws = target.eval(xs, log=True) - proposal.eval(xs, log=True)
    log_ws -= scipy.misc.logsumexp(log_ws)

    return xs, log_ws
//  [python]
def importance_sample(target, proposal, n_samples, rng=np.random):
    """
    Importance sampling.
    :param target: target distribution
    :param proposal: proposal distribution
    :param n_samples: number of samples
    :param rng: random generator to use
    :return: samples, normalized log weights
    """

    xs = proposal.gen(n_samples, rng=rng)
    log_ws = target.eval(xs, log=True) - proposal.eval(xs, log=True)
    log_ws -= scipy.special.logsumexp(log_ws)

    return xs, log_ws
//  [python]
def test_logsumexp_single_infinite():
    a = np.inf
    ans = pymbar.utils.logsumexp(a)
    ans_scipy = logsumexp(a)
    assert_equal(ans, ans_scipy)
//  [python]
def test_logsumexp_single_infinite():
    a = np.inf
    ans = pymbar.utils.logsumexp(a)
    ans_scipy = scipy.special.logsumexp(a)
    assert_equal(ans, ans_scipy)
//  [python]
def posterior_predictive(C, Y, K, alpha, beta, gamma):
    """
    computes p(y | Y,C) for all possible 2^D possible values
    """

    N, D = Y.shape
    assert C.shape[0] == N

    nks = np.bincount(C, minlength=K)
    sums = np.zeros((K, D))
    for yi, ci in zip(Y, C):
        sums[ci] += yi

    def fn(yvalue):
        def fn1(nk, sum_yid, yd):
            assert nk >= sum_yid
            theta = (beta + sum_yid) / (beta + gamma + nk)
            assert theta >= 0.0 and theta <= 1.0
            return np.log(theta) if yd else np.log(1.-theta)
        def fn2(nk, row):
            assert len(yvalue) == row.shape[0]
            term1 = np.log(nk + alpha/K) - np.log(N + alpha)
            term2 = sum(fn1(nk, sum_yid, yd) for sum_yid, yd in zip(row, yvalue))
            return term1 + term2
        return sp.misc.logsumexp([fn2(nk, row) for nk, row in zip(nks, sums)])

    yvalues = it.product([0, 1], repeat=D)
    lg_pr_yvalue = map(fn, yvalues)
    return np.exp(lg_pr_yvalue)
//  [python]
def posterior_predictive(C, Y, K, alpha, beta, gamma):
    """
    computes p(y | Y,C) for all possible 2^D possible values
    """

    N, D = Y.shape
    assert C.shape[0] == N

    nks = np.bincount(C, minlength=K)
    sums = np.zeros((K, D))
    for yi, ci in zip(Y, C):
        sums[ci] += yi

    def fn(yvalue):
        def fn1(nk, sum_yid, yd):
            assert nk >= sum_yid
            theta = (beta + sum_yid) / (beta + gamma + nk)
            assert theta >= 0.0 and theta <= 1.0
            return np.log(theta) if yd else np.log(1.-theta)
        def fn2(nk, row):
            assert len(yvalue) == row.shape[0]
            term1 = np.log(nk + alpha/K) - np.log(N + alpha)
            term2 = sum(fn1(nk, sum_yid, yd) for sum_yid, yd in zip(row, yvalue))
            return term1 + term2
        return sp.special.logsumexp([fn2(nk, row) for nk, row in zip(nks, sums)])

    yvalues = it.product([0, 1], repeat=D)
    lg_pr_yvalue = map(fn, yvalues)
    return np.exp(lg_pr_yvalue)
//  [python]
def soft_sample(self, evidence_data, rand_gen=None):

        if rand_gen is None:
            rand_gen = np.random.RandomState(RAND_SEED)

        n_samples = evidence_data.shape[0]
        n_features = evidence_data.shape[1]
        n_children = len(self.children)

        sample_children_evals = [c.sample(evidence_data, rand_gen) for c in self.children]
        children_log_probs = np.array([c_lp for c_lp, _c_r in sample_children_evals])
        children_samples = np.array([c_r for _c_lp, c_r in sample_children_evals])

        w_children_log_probs = self.log_weights + children_log_probs

        children_probs = np.exp(w_children_log_probs)
        children_probs = children_probs / children_probs.sum(axis=0)
        sample_log_probs = scipy.misc.logsumexp(w_children_log_probs, axis=0)

        samples = children_probs[:, :, None] * children_samples
        samples = np.sum(samples, axis=0)

        return sample_log_probs, samples
//  [python]
def soft_sample(self, evidence_data, rand_gen=None):

        if rand_gen is None:
            rand_gen = np.random.RandomState(RAND_SEED)

        n_samples = evidence_data.shape[0]
        n_features = evidence_data.shape[1]
        n_children = len(self.children)

        sample_children_evals = [c.sample(evidence_data, rand_gen) for c in self.children]
        children_log_probs = np.array([c_lp for c_lp, _c_r in sample_children_evals])
        children_samples = np.array([c_r for _c_lp, c_r in sample_children_evals])

        w_children_log_probs = self.log_weights + children_log_probs

        children_probs = np.exp(w_children_log_probs)
        children_probs = children_probs / children_probs.sum(axis=0)
        sample_log_probs = scipy.special.logsumexp(w_children_log_probs, axis=0)

        samples = children_probs[:, :, None] * children_samples
        samples = np.sum(samples, axis=0)

        return sample_log_probs, samples
//  [python]
def test_logsumexp():
    a = np.random.normal(size=(200, 500, 5))

    for axis in range(a.ndim):
        ans_ne = pymbar.utils.logsumexp(a, axis=axis)
        ans_no_ne = pymbar.utils.logsumexp(a, axis=axis, use_numexpr=False)
        ans_scipy = logsumexp(a, axis=axis)
        assert_almost_equal(ans_ne, ans_no_ne)
        assert_almost_equal(ans_ne, ans_scipy)
//  [python]
def test_logsumexp():
    a = np.random.normal(size=(200, 500, 5))

    for axis in range(a.ndim):
        ans_ne = pymbar.utils.logsumexp(a, axis=axis)
        ans_no_ne = pymbar.utils.logsumexp(a, axis=axis, use_numexpr=False)
        ans_scipy = scipy.special.logsumexp(a, axis=axis)
        assert_almost_equal(ans_ne, ans_no_ne)
        assert_almost_equal(ans_ne, ans_scipy)
//  [python]
def test_logsumexp_b():
    a = np.random.normal(size=(200, 500, 5))
    b = np.random.normal(size=(200, 500, 5)) ** 2.0

    for axis in range(a.ndim):
        ans_ne = pymbar.utils.logsumexp(a, b=b, axis=axis)
        ans_no_ne = pymbar.utils.logsumexp(a, b=b, axis=axis, use_numexpr=False)
        ans_scipy = logsumexp(a, b=b, axis=axis)
        assert_almost_equal(ans_ne, ans_no_ne)
        assert_almost_equal(ans_ne, ans_scipy)
//  [python]
def test_logsumexp_b():
    a = np.random.normal(size=(200, 500, 5))
    b = np.random.normal(size=(200, 500, 5)) ** 2.0

    for axis in range(a.ndim):
        ans_ne = pymbar.utils.logsumexp(a, b=b, axis=axis)
        ans_no_ne = pymbar.utils.logsumexp(a, b=b, axis=axis, use_numexpr=False)
        ans_scipy = scipy.special.logsumexp(a, b=b, axis=axis)
        assert_almost_equal(ans_ne, ans_no_ne)
        assert_almost_equal(ans_ne, ans_scipy)
//  [python]
def xalphas(ms):
    mslen = len(ms)
    alpha = np.empty((mslen, ms[0].shape[0]))
    alpha[0,:] = ms[0][0,:]
    for i in range(1,mslen-1):
        alpha[i,:] = [ misc.logsumexp(alpha[i-1,:] + ms[i][:,y1]) for y1 in range(ms[i].shape[1]) ]
    alpha[-1,:] = alpha[-2,:] + ms[-1][:,0]
    return alpha
//  [python]
def xalphas(ms):
    mslen = len(ms)
    alpha = np.empty((mslen, ms[0].shape[0]))
    alpha[0,:] = ms[0][0,:]
    for i in range(1,mslen-1):
        alpha[i,:] = [ scipy.special.logsumexp(alpha[i-1,:] + ms[i][:,y1]) for y1 in range(ms[i].shape[1]) ]
    alpha[-1,:] = alpha[-2,:] + ms[-1][:,0]
    return alpha
//  [python]
def xbetas(ms):
    mslen = len(ms)
    beta = np.zeros((mslen, ms[0].shape[0]))
    beta[-1,:] = ms[-1][:,0]
    for i in range(mslen-2, 0, -1):
        beta[i,:] = [ misc.logsumexp(ms[i][y0,:] + beta[i+1,:]) for y0 in range(ms[i].shape[0]) ]
    beta[0,:] = ms[0][0,:] + beta[1,:]
    return beta
//  [python]
def xbetas(ms):
    mslen = len(ms)
    beta = np.zeros((mslen, ms[0].shape[0]))
    beta[-1,:] = ms[-1][:,0]
    for i in range(mslen-2, 0, -1):
        beta[i,:] = [ scipy.special.logsumexp(ms[i][y0,:] + beta[i+1,:]) for y0 in range(ms[i].shape[0]) ]
    beta[0,:] = ms[0][0,:] + beta[1,:]
    return beta
//  [python]
Ls = []
while True:
	N = 1000000
	x = numpy.random.uniform(size=(N, ndim))
	L = loglikelihood(x)
	Ls.append(scipy.misc.logsumexp(L) - log(N))
	means = []
	for i in range(20):
		vals = numpy.exp([Ls[j] for j in numpy.random.randint(0, len(Ls), size=len(Ls))])
		means.append(numpy.log(vals.mean()))
	print('%.5f BS: %.5f +- %.5f (%d x %d)' % (numpy.log(numpy.exp(Ls).mean()), numpy.mean(means), numpy.std(means), len(Ls), N))
//  [python]
Ls = []
while True:
	N = 1000000
	x = numpy.random.uniform(size=(N, ndim))
	L = loglikelihood(x)
	Ls.append(scipy.special.logsumexp(L) - log(N))
	means = []
	for i in range(20):
		vals = numpy.exp([Ls[j] for j in numpy.random.randint(0, len(Ls), size=len(Ls))])
		means.append(numpy.log(vals.mean()))
	print('%.5f BS: %.5f +- %.5f (%d x %d)' % (numpy.log(numpy.exp(Ls).mean()), numpy.mean(means), numpy.std(means), len(Ls), N))
//  [python]
def bound(self, corpus, gamma=None, subsample_ratio=1.0):
        """
        Estimate the variational bound of documents from `corpus`:
        E_q[log p(corpus)] - E_q[log q(corpus)]

        `gamma` are the variational parameters on topic weights for each `corpus`
        document (=2d matrix=what comes out of `inference()`).
        If not supplied, will be inferred from the model.

        """
        score = 0.0
        _lambda = self.state.get_lambda()
        Elogbeta = dirichlet_expectation(_lambda)

        for d, doc in enumerate(corpus):  # stream the input doc-by-doc, in case it's too large to fit in RAM
            if d % self.chunksize == 0:
                logger.debug("bound: at document #%i" % d)
            if gamma is None:
                gammad, _ = self.inference([doc])
            else:
                gammad = gamma[d]
            Elogthetad = dirichlet_expectation(gammad)

            # E[log p(doc | theta, beta)]
            score += numpy.sum(cnt * logsumexp(Elogthetad + Elogbeta[:, id]) for id, cnt in doc)

            # E[log p(theta | alpha) - log q(theta | gamma)]; assumes alpha is a vector
            score += numpy.sum((self.alpha - gammad) * Elogthetad)
            score += numpy.sum(gammaln(gammad) - gammaln(self.alpha))
            score += gammaln(numpy.sum(self.alpha)) - gammaln(numpy.sum(gammad))

        # compensate likelihood for when `corpus` above is only a sample of the whole corpus
        score *= subsample_ratio

        # E[log p(beta | eta) - log q (beta | lambda)]; assumes eta is a scalar
        score += numpy.sum((self.eta - _lambda) * Elogbeta)
        score += numpy.sum(gammaln(_lambda) - gammaln(self.eta))
        score += numpy.sum(gammaln(self.eta * self.num_terms) - gammaln(numpy.sum(_lambda, 1)))
        return score
//  [python]
def bound(self, corpus, gamma=None, subsample_ratio=1.0):
        """
        Estimate the variational bound of documents from `corpus`:
        E_q[log p(corpus)] - E_q[log q(corpus)]

        `gamma` are the variational parameters on topic weights for each `corpus`
        document (=2d matrix=what comes out of `inference()`).
        If not supplied, will be inferred from the model.

        """
        score = 0.0
        _lambda = self.state.get_lambda()
        Elogbeta = dirichlet_expectation(_lambda)

        for d, doc in enumerate(corpus):  # stream the input doc-by-doc, in case it's too large to fit in RAM
            if d % self.chunksize == 0:
                logger.debug("bound: at document #%i" % d)
            if gamma is None:
                gammad, _ = self.inference([doc])
            else:
                gammad = gamma[d]
            Elogthetad = dirichlet_expectation(gammad)

            # E[log p(doc | theta, beta)]
            score += numpy.sum(cnt * scipy.special.logsumexp(Elogthetad + Elogbeta[:, id]) for id, cnt in doc)

            # E[log p(theta | alpha) - log q(theta | gamma)]; assumes alpha is a vector
            score += numpy.sum((self.alpha - gammad) * Elogthetad)
            score += numpy.sum(gammaln(gammad) - gammaln(self.alpha))
            score += gammaln(numpy.sum(self.alpha)) - gammaln(numpy.sum(gammad))

        # compensate likelihood for when `corpus` above is only a sample of the whole corpus
        score *= subsample_ratio

        # E[log p(beta | eta) - log q (beta | lambda)]; assumes eta is a scalar
        score += numpy.sum((self.eta - _lambda) * Elogbeta)
        score += numpy.sum(gammaln(_lambda) - gammaln(self.eta))
        score += numpy.sum(gammaln(self.eta * self.num_terms) - gammaln(numpy.sum(_lambda, 1)))
        return score
//  [python]
def model_values(weights, yset, features, data_x):
  return crfext.model_values(list(weights), yset, features, data_x)
  val = 0.
  for x, p in data_x:
    val += p * misc.logsumexp( xalphas( xmatrices(yset, features, weights, x ) )[-1])
  return val
//  [python]
def model_values(weights, yset, features, data_x):
  return crfext.model_values(list(weights), yset, features, data_x)
  val = 0.
  for x, p in data_x:
    val += p * scipy.special.logsumexp( xalphas( xmatrices(yset, features, weights, x ) )[-1])
  return val
//  [python]
def log_partition(alpha,f_sim,f_exp):
    """Return the log partition function of an alpha-reweighted ensemble."""
    q = -1*f_sim.dot(alpha).astype('float128')
    return scipy.misc.logsumexp(q)
//  [python]
def log_partition(alpha,f_sim,f_exp):
    """Return the log partition function of an alpha-reweighted ensemble."""
    q = -1*f_sim.dot(alpha).astype('float128')
    return scipy.special.logsumexp(q)
//  [python]
def log_softmax(xs):
  assert len(xs.shape) == 2
  import scipy.misc
  return xs - scipy.misc.logsumexp(xs, 1)[:, None]
//  [python]
def log_softmax(xs):
  assert len(xs.shape) == 2
  import scipy.special
  return xs - scipy.special.logsumexp(xs, 1)[:, None]
//  [python]
def model_grads(yset, features, weights, data_x):
  return np.asarray( crfext.model_grads(yset, features, list(weights), data_x) )
  K = len(weights)
  grads = np.zeros( K )
  malphabetaz = []
  for x, p in data_x:
    ms = xmatrices(yset, features, weights, x)
    alpha = xalphas(ms)
    beta = xbetas(ms)
    logz1 = misc.logsumexp( alpha[-1] )
    logz2 = misc.logsumexp( beta[0] )
    assert np.allclose(logz1, logz2)
    malphabetaz.append( (x, p, ms, alpha, beta, logz1) )
  for k in range(K):
    for x, p, ms, alpha, beta, logz in malphabetaz:
      grads[k] += p * kx(yset, features[k], x, ms, alpha, beta, logz)
  return grads
//  [python]
def model_grads(yset, features, weights, data_x):
  return np.asarray( crfext.model_grads(yset, features, list(weights), data_x) )
  K = len(weights)
  grads = np.zeros( K )
  malphabetaz = []
  for x, p in data_x:
    ms = xmatrices(yset, features, weights, x)
    alpha = xalphas(ms)
    beta = xbetas(ms)
    logz1 = scipy.special.logsumexp( alpha[-1] )
    logz2 = scipy.special.logsumexp( beta[0] )
    assert np.allclose(logz1, logz2)
    malphabetaz.append( (x, p, ms, alpha, beta, logz1) )
  for k in range(K):
    for x, p, ms, alpha, beta, logz in malphabetaz:
      grads[k] += p * kx(yset, features[k], x, ms, alpha, beta, logz)
  return grads
//  [python]
def lognormexp(values, dim=0):
    """Log of exponentiates and normalizes a Tensor/Variable/np.ndarray.

    input:
        values: Tensor/Variable/np.ndarray [dim_1, ..., dim_N]
        dim: n
    output:
        result: Tensor/Variable/np.ndarray [dim_1, ..., dim_N]
            where result[i_1, ..., i_N] =

                            exp(values[i_1, ..., i_N])
            ------------------------------------------------------------
             sum_{j = 1}^{dim_n} exp(values[i_1, ..., j, ..., i_N])
    """

    if isinstance(values, np.ndarray):
        log_denominator = scipy.misc.logsumexp(
            values, axis=dim, keepdims=True
        )
        # log_numerator = values
        return values - log_denominator
    else:
        log_denominator = logsumexp(values, dim=dim, keepdim=True)
        # log_numerator = values
        return values - log_denominator.expand_as(values)
//  [python]
def lognormexp(values, dim=0):
    """Log of exponentiates and normalizes a Tensor/Variable/np.ndarray.

    input:
        values: Tensor/Variable/np.ndarray [dim_1, ..., dim_N]
        dim: n
    output:
        result: Tensor/Variable/np.ndarray [dim_1, ..., dim_N]
            where result[i_1, ..., i_N] =

                            exp(values[i_1, ..., i_N])
            ------------------------------------------------------------
             sum_{j = 1}^{dim_n} exp(values[i_1, ..., j, ..., i_N])
    """

    if isinstance(values, np.ndarray):
        log_denominator = scipy.special.logsumexp(
            values, axis=dim, keepdims=True
        )
        # log_numerator = values
        return values - log_denominator
    else:
        log_denominator = logsumexp(values, dim=dim, keepdim=True)
        # log_numerator = values
        return values - log_denominator.expand_as(values)
//  [python]
def model_grads1(yset, features, weights, data_x):
  return np.asarray( crfext.model_grads1(yset, features, list(weights), data_x) )
  K = len(weights)
  grads = np.zeros( K )
  malphabetaz = []
  for x, p in data_x:
    ms = xmatrices(yset, features, weights, x)
    alpha = xalphas(ms)
    beta = xbetas(ms)
    logz1 = misc.logsumexp( alpha[-1] )
    logz2 = misc.logsumexp( beta[0] )
    assert np.allclose(logz1, logz2)
    malphabetaz.append( (x, p, ms, alpha, beta, logz1) )
  global kxtablecache
  if not kxtablecache: kxtablecache = create_kxcache(yset, features, data_x)
  for k in range(K):
    for xnum, xs0s1i in kxtablecache[k]:
      _, proba, ms, alpha, beta, logz = malphabetaz[xnum]
      grads[k] += proba * np.sum( [ xyiiproba(yset, ms, alpha, beta, s0, s1, i, logz) for s0, s1, i in xs0s1i ] )
  return grads
//  [python]
def model_grads1(yset, features, weights, data_x):
  return np.asarray( crfext.model_grads1(yset, features, list(weights), data_x) )
  K = len(weights)
  grads = np.zeros( K )
  malphabetaz = []
  for x, p in data_x:
    ms = xmatrices(yset, features, weights, x)
    alpha = xalphas(ms)
    beta = xbetas(ms)
    logz1 = scipy.special.logsumexp( alpha[-1] )
    logz2 = scipy.special.logsumexp( beta[0] )
    assert np.allclose(logz1, logz2)
    malphabetaz.append( (x, p, ms, alpha, beta, logz1) )
  global kxtablecache
  if not kxtablecache: kxtablecache = create_kxcache(yset, features, data_x)
  for k in range(K):
    for xnum, xs0s1i in kxtablecache[k]:
      _, proba, ms, alpha, beta, logz = malphabetaz[xnum]
      grads[k] += proba * np.sum( [ xyiiproba(yset, ms, alpha, beta, s0, s1, i, logz) for s0, s1, i in xs0s1i ] )
  return grads
//  [python]
def evaluate_logprob(model, split, use_image_space=False, return_avg=True, batch=2000):
    """
    Evaluate a trained model only in terms of log probability.
    :param model: the model to evaluate. Can be any made, maf, or real nvp
    :param split: string, the data split to evaluate on. Must be 'trn', 'val' or 'tst'
    :param use_image_space: bool, whether to report log probability in [0, 1] image space (only for cifar and mnist)
    :param return_avg: bool, whether to return average log prob with std error, or all log probs
    :param batch: batch size to use for computing log probability
    :return: average log probability & standard error, or all log probs
    """

    assert is_data_loaded(), 'Dataset hasn\'t been loaded'

    # choose which data split to evaluate on
    data_split = getattr(data, split, None)
    if data_split is None:
        raise ValueError('Invalid data split')

    if is_conditional(model):

        logprobs = np.empty([data_split.N, data.n_labels])

        for i in range(data.n_labels):

            # create labels
            y = np.zeros([data_split.N, data.n_labels])
            y[:, i] = 1

            # process data in batches to make sure they fit in memory
            r, l = 0, batch
            while r < data_split.N:
                logprobs[r:l, i] = model.eval([y[r:l], data_split.x[r:l]])
                l += batch
                r += batch

        logprobs = scipy.misc.logsumexp(logprobs, axis=1) - np.log(logprobs.shape[1])

    else:

        logprobs = np.empty(data_split.N)

        # process data in batches to make sure they fit in memory
        r, l = 0, batch
        while r < data_split.N:
            logprobs[r:l] = model.eval(data_split.x[r:l])
            l += batch
            r += batch

    if use_image_space:
        assert data_name in ['mnist', 'cifar10']
        z = util.logistic(data_split.x)
        logprobs += data.n_dims * np.log(1-2*data.alpha) - np.sum(np.log(z) + np.log(1-z), axis=1)

    if return_avg:
        avg_logprob = logprobs.mean()
        std_err = logprobs.std() / np.sqrt(data_split.N)
        return avg_logprob, std_err

    else:
        return logprobs