jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def apply_mask(real_obs, mask, means, stds, mix_coefs=None, delta=0.1, encoder=None):
  if mix_coefs is None:
    mix_coefs = np.ones(1)
    means = means[np.newaxis]
    stds = stds[np.newaxis]

  all_idxes = np.arange(0, real_obs.shape[1], 1)
  masked_obs = deepcopy(real_obs)
  kldivs = np.ones(real_obs.shape[0]) * np.nan
  mix_ids = list(range(len(mix_coefs)))
  log_mix_coefs = np.log(mix_coefs)
  for i in range(real_obs.shape[0]):
    mask_mean = mask[i, :].mean()
    idxes1 = all_idxes[mask[i, :] < 1]
    idxes2 = all_idxes[mask[i, :] > 0]

    if mask_mean == 0:
      log_mix_ps = log_mix_coefs
    else:
      log_mix_ps = [mvn_logpdf(
        means[j, idxes2],
        stds[j, idxes2],
        masked_obs[i, idxes2],
        delta=delta
        )+log_mix_coefs[j] for j in range(means.shape[0])]
      log_mix_ps = np.maximum(-1e6, log_mix_ps)
      log_mix_ps -= np.max(log_mix_ps)
      log_mix_ps = np.log(np.sum(np.exp(log_mix_ps)))
    mix_id = np.random.choice(mix_ids, p=np.exp(log_mix_ps))

    mean = means[mix_id]
    std = stds[mix_id]

    if encoder is None:
      mu1 = mean[idxes1]
      std1 = std[idxes1]
      rand_obs = np.random.normal(mu1, std1)
    else:
      rand_obs = encoder.sample(1)[0, idxes1]
    m = (np.random.random(len(idxes1)) < mask[i, idxes1]).astype(float)
    masked_obs[i, idxes1] = m * masked_obs[i, idxes1] + (1 - m) * rand_obs

    mu2 = mean[idxes2]
    s22_diag = std[idxes2]
    masked_obs[i, idxes2] = disc_norm_rv(masked_obs[i, idxes2], mu2, s22_diag, delta=delta)
    p = mvn_logpdf(mu2, s22_diag, masked_obs[i, idxes2], delta=delta)
    # bits back
    kldivs[i] = log_mix_ps[mix_id] - log_mix_coefs[mix_id] - p
    kldivs[i] /= np.log(2)
  return masked_obs, kldivs
//  [python]
def apply_mask(real_obs, mask, means, stds, mix_coefs=None, delta=0.1, encoder=None):
  if mix_coefs is None:
    mix_coefs = np.ones(1)
    means = means[np.newaxis]
    stds = stds[np.newaxis]

  all_idxes = np.arange(0, real_obs.shape[1], 1)
  masked_obs = deepcopy(real_obs)
  kldivs = np.ones(real_obs.shape[0]) * np.nan
  mix_ids = list(range(len(mix_coefs)))
  log_mix_coefs = np.log(mix_coefs)
  for i in range(real_obs.shape[0]):
    mask_mean = mask[i, :].mean()
    idxes1 = all_idxes[mask[i, :] < 1]
    idxes2 = all_idxes[mask[i, :] > 0]

    if mask_mean == 0:
      log_mix_ps = log_mix_coefs
    else:
      log_mix_ps = [mvn_logpdf(
        means[j, idxes2],
        stds[j, idxes2],
        masked_obs[i, idxes2],
        delta=delta
        )+log_mix_coefs[j] for j in range(means.shape[0])]
      log_mix_ps = np.maximum(-1e6, log_mix_ps)
      log_mix_ps -= scipy.special.logsumexp(log_mix_ps)
    mix_id = np.random.choice(mix_ids, p=np.exp(log_mix_ps))

    mean = means[mix_id]
    std = stds[mix_id]

    if encoder is None:
      mu1 = mean[idxes1]
      std1 = std[idxes1]
      rand_obs = np.random.normal(mu1, std1)
    else:
      rand_obs = encoder.sample(1)[0, idxes1]
    m = (np.random.random(len(idxes1)) < mask[i, idxes1]).astype(float)
    masked_obs[i, idxes1] = m * masked_obs[i, idxes1] + (1 - m) * rand_obs

    mu2 = mean[idxes2]
    s22_diag = std[idxes2]
    masked_obs[i, idxes2] = disc_norm_rv(masked_obs[i, idxes2], mu2, s22_diag, delta=delta)
    p = mvn_logpdf(mu2, s22_diag, masked_obs[i, idxes2], delta=delta)
    # bits back
    kldivs[i] = log_mix_ps[mix_id] - log_mix_coefs[mix_id] - p
    kldivs[i] /= np.log(2)
  return masked_obs, kldivs
//  [python]
def test_segmented_cuda_logsumexp_edge_cases(precision, atol, rtol, seed, count, edge_case):
    max_values_per_segment = 20
    rng = np.random.default_rng(seed)
    # Shuffle the ordering of the edge case, should be independent of ordering
    rng.shuffle(edge_case)
    # Only one sample will
    samples = [edge_case * count, rng.normal(size=max_values_per_segment).tolist()]

    rng.shuffle(samples)

    if precision == np.float32:
        summer = custom_ops.SegmentedSumExp_f32(max_values_per_segment, len(samples))
    else:
        summer = custom_ops.SegmentedSumExp_f64(max_values_per_segment, len(samples))

    test_vals = summer.logsumexp(samples)
    assert len(test_vals) == len(samples)
    # Verify that the results are deterministic
    np.testing.assert_array_equal(test_vals, summer.logsumexp(samples))

    for test_val, vals in zip(test_vals, samples):
        ref_logsumexp = s.logsumexp(vals)
        if np.isfinite(ref_logsumexp):
            np.testing.assert_allclose(ref_logsumexp, test_val, atol=atol, rtol=rtol)
        else:
            assert np.isnan(test_val)
//  [python]
def test_segmented_cuda_logsumexp_edge_cases(precision, atol, rtol, seed, count, edge_case):
    max_values_per_segment = 20
    rng = np.random.default_rng(seed)
    # Shuffle the ordering of the edge case, should be independent of ordering
    rng.shuffle(edge_case)
    # Only one sample will
    samples = [edge_case * count, rng.normal(size=max_values_per_segment).tolist()]

    rng.shuffle(samples)

    if precision == np.float32:
        summer = custom_ops.SegmentedSumExp_f32(max_values_per_segment, len(samples))
    else:
        summer = custom_ops.SegmentedSumExp_f64(max_values_per_segment, len(samples))

    test_vals = summer.logsumexp(samples)
    assert len(test_vals) == len(samples)
    # Verify that the results are deterministic
    np.testing.assert_array_equal(test_vals, summer.logsumexp(samples))

    for test_val, vals in zip(test_vals, samples):
        ref_logsumexp = logsumexp(vals)
        if np.isfinite(ref_logsumexp):
            np.testing.assert_allclose(ref_logsumexp, test_val, atol=atol, rtol=rtol)
        else:
            assert np.isnan(test_val)
//  [python]
def test_logsumexp_numba(self):
        """helper.numba.logsumexp"""
        from scipy.special import logsumexp as s
        a = np.arange(100).reshape((2, -1))
        assert numba.logsumexp(a) == pytest.approx(s(a))
        assert numba.logsumexp(np.zeros(1)) == pytest.approx(
            s(np.zeros(1)))
//  [python]
def test_logsumexp_numba(self):
        """helper.numba.logsumexp"""
        from scipy.special import logsumexp
        a = np.arange(100).reshape((2, -1))
        assert numba.logsumexp(a) == pytest.approx(logsumexp(a))
        assert numba.logsumexp(np.zeros(1)) == pytest.approx(
            logsumexp(np.zeros(1)))
//  [python]
def logsumexp2(z):
    """Multiprocessing pool.map requires a top-level function."""
    return logsumexp(z, axis=2)
//  [python]
def calculate_pac_dir(prob_np_all, label_np_all, alpha):
  """Compute the PACTran-Dirichlet estimator."""
  starttime = time.time()
  label_np_all = one_hot(label_np_all)  # [n, v]
  soft_targets_sum = np.sum(label_np_all, axis=0)  # [v]
  soft_targets_sum = np.expand_dims(soft_targets_sum, axis=1)  # [v, 1]
  a0 = alpha * soft_targets_sum / np.sum(soft_targets_sum) + 1e-10

  # initialize
  qz = prob_np_all  # [n, d]
  log_s = np.log(prob_np_all + 1e-10)  # [n, d]

  for _ in range(10):
    aw = a0 + np.sum(np.einsum("BY,BZ->BYZ", label_np_all, qz), axis=0)
    logits_qz = (log_s +
                 np.matmul(label_np_all, scipy.special.digamma(aw)) -
                 np.reshape(scipy.special.digamma(np.sum(aw, axis=0)), [1, -1]))
    log_qz = logits_qz - np.max(logits_qz, axis=-1, keepdims=True)
    log_qz = log_qz - np.log(np.sum(np.exp(log_qz), axis=-1, keepdims=True))
    qz = np.exp(log_qz)

  log_c0 = scipy.special.loggamma(np.sum(a0)) - np.sum(
      scipy.special.loggamma(a0))
  log_c = scipy.special.loggamma(np.sum(aw, axis=0)) - np.sum(
      scipy.special.loggamma(aw), axis=0)

  pac_dir = np.sum(
      log_c0 - log_c - np.sum(qz * (log_qz - log_s), axis=0))
  pac_dir = -pac_dir / label_np_all.size
  endtime = time.time()
  return [("pac_dir_%.1f" % alpha, pac_dir),
          ("time", endtime - starttime)]
//  [python]
def calculate_pac_dir(prob_np_all, label_np_all, alpha):
  """Compute the PACTran-Dirichlet estimator."""
  starttime = time.time()
  label_np_all = one_hot(label_np_all)  # [n, v]
  soft_targets_sum = np.sum(label_np_all, axis=0)  # [v]
  soft_targets_sum = np.expand_dims(soft_targets_sum, axis=1)  # [v, 1]
  a0 = alpha * soft_targets_sum / np.sum(soft_targets_sum) + 1e-10

  # initialize
  qz = prob_np_all  # [n, d]
  log_s = np.log(prob_np_all + 1e-10)  # [n, d]

  for _ in range(10):
    aw = a0 + np.sum(np.einsum("BY,BZ->BYZ", label_np_all, qz), axis=0)
    logits_qz = (log_s +
                 np.matmul(label_np_all, scipy.special.digamma(aw)) -
                 np.reshape(scipy.special.digamma(np.sum(aw, axis=0)), [1, -1]))
    log_qz = logits_qz - scipy.special.logsumexp(
        logits_qz, axis=-1, keepdims=True)
    qz = np.exp(log_qz)

  log_c0 = scipy.special.loggamma(np.sum(a0)) - np.sum(
      scipy.special.loggamma(a0))
  log_c = scipy.special.loggamma(np.sum(aw, axis=0)) - np.sum(
      scipy.special.loggamma(aw), axis=0)

  pac_dir = np.sum(
      log_c0 - log_c - np.sum(qz * (log_qz - log_s), axis=0))
  pac_dir = -pac_dir / label_np_all.size
  endtime = time.time()
  return [("pac_dir_%.1f" % alpha, pac_dir),
          ("time", endtime - starttime)]
//  [python]
def normalize_latent(self, log_p_y_given_x_unnorm):
        """Normalize the latent variable distribution

        For each sample in the training set, we estimate a probability distribution
        over y_j, each hidden factor. Here we normalize it. (Eq. 7 in paper.)
        This normalization factor is quite useful as described in upcoming work.

        Parameters
        ----------
        Unnormalized distribution of hidden factors for each training sample.

        Returns
        -------
        p_y_given_x : 3D array, shape (n_hidden, n_samples, dim_hidden)
            p(y_j|x^l), the probability distribution over all hidden factors,
            for data samples l = 1...n_samples
        log_z : 2D array, shape (n_hidden, n_samples)
            Point-wise estimate of total correlation explained by each Y_j for each sample,
            used to estimate overall total correlation.

        """

        log_z = np.logaddexp.reduce(log_p_y_given_x_unnorm, axis=2)  # Essential to maintain precision.
        log_z = log_z.reshape((self.n_hidden, -1, 1))

        return np.exp(log_p_y_given_x_unnorm - log_z), log_z
//  [python]
def normalize_latent(self, log_p_y_given_x_unnorm):
        """Normalize the latent variable distribution

        For each sample in the training set, we estimate a probability distribution
        over y_j, each hidden factor. Here we normalize it. (Eq. 7 in paper.)
        This normalization factor is quite useful as described in upcoming work.

        Parameters
        ----------
        Unnormalized distribution of hidden factors for each training sample.

        Returns
        -------
        p_y_given_x : 3D array, shape (n_hidden, n_samples, dim_hidden)
            p(y_j|x^l), the probability distribution over all hidden factors,
            for data samples l = 1...n_samples
        log_z : 2D array, shape (n_hidden, n_samples)
            Point-wise estimate of total correlation explained by each Y_j for each sample,
            used to estimate overall total correlation.

        """

        log_z = logsumexp(log_p_y_given_x_unnorm, axis=2)  # Essential to maintain precision.
        log_z = log_z.reshape((self.n_hidden, -1, 1))

        return np.exp(log_p_y_given_x_unnorm - log_z), log_z
//  [python]
def _get_delta_error_term(
        Lxs: typing.Sequence[float],
        ps: typing.Sequence[float],
        num_compositions: int = 500,
        L: float = 20.0,
        lambd: typing.Optional[float] = None
    ) -> float:
    """ Computes the total error term for δ computed by the Fourier accountant
    for repeated application of a privacy mechanism.

    The computation follows Theorem 7 in Koskela & Honkela, "Computing Differential Privacy for
    Heterogeneous Compositions Using FFT", 2021, arXiv preprint, https://arxiv.org/abs/2102.12412 .

    Args:
        - Lxs: Sequence of privacy loss values.
        - ps: Sequence of privacy loss probability masses.
        - num_compositions: The number of compositions (=applications) of the privacy mechanism.
        - L: The truncation threshold (in privacy loss space) used by the accountant.
        - lambd: The parameter λ for error estimation.
    """

    if lambd is None:
        lambd = .5 * L

    assert np.size(ps) == np.size(Lxs)
    nonzero_probability_filter = ~np.isclose(ps, 0)
    ps = ps[nonzero_probability_filter]
    Lxs = Lxs[nonzero_probability_filter]
    assert np.all(ps > 0)

    # Compute the lambda-divergence \alpha^+
    alpha_plus = np.max(np.log(ps) + lambd * Lxs) + np.log(np.sum(np.exp(np.log(ps) + lambd * Lxs - np.max(np.log(ps) + lambd * Lxs))))

    # Compute the lambda-divergence \alpha^-
    alpha_minus = np.max(np.log(ps) - lambd * Lxs) + np.log(np.sum(np.exp(np.log(ps) - lambd * Lxs - np.max(np.log(ps) - lambd * Lxs))))

    k = num_compositions

    common_factor_log = -(L * lambd + np.log1p(-np.exp(-2 * L * lambd)))

    T1_log = k * alpha_plus + common_factor_log
    T2_log = k * alpha_minus + common_factor_log

    T_max_log = np.maximum(T1_log, T2_log)

    error_term = np.exp(T_max_log) * (np.exp(T1_log - T_max_log) + np.exp(T2_log - T_max_log))

    return error_term
//  [python]
def _get_delta_error_term(
        Lxs: typing.Sequence[float],
        ps: typing.Sequence[float],
        num_compositions: int = 500,
        L: float = 20.0,
        lambd: typing.Optional[float] = None
    ) -> float:
    """ Computes the total error term for δ computed by the Fourier accountant
    for repeated application of a privacy mechanism.

    The computation follows Theorem 7 in Koskela & Honkela, "Computing Differential Privacy for
    Heterogeneous Compositions Using FFT", 2021, arXiv preprint, https://arxiv.org/abs/2102.12412 .

    Args:
        - Lxs: Sequence of privacy loss values.
        - ps: Sequence of privacy loss probability masses.
        - num_compositions: The number of compositions (=applications) of the privacy mechanism.
        - L: The truncation threshold (in privacy loss space) used by the accountant.
        - lambd: The parameter λ for error estimation.
    """

    if lambd is None:
        lambd = .5 * L

    assert np.size(ps) == np.size(Lxs)
    nonzero_probability_filter = ~np.isclose(ps, 0)
    ps = ps[nonzero_probability_filter]
    Lxs = Lxs[nonzero_probability_filter]
    assert np.all(ps > 0)

    # Compute the lambda-divergence \alpha^+
    alpha_plus = scipy.special.logsumexp(np.log(ps) + lambd * Lxs)

    # Compute the lambda-divergence \alpha^-
    alpha_minus = scipy.special.logsumexp(np.log(ps) - lambd * Lxs)

    k = num_compositions

    common_factor_log = -(L * lambd + np.log1p(-np.exp(-2 * L * lambd)))

    T1_log = k * alpha_plus + common_factor_log
    T2_log = k * alpha_minus + common_factor_log

    T_max_log = np.maximum(T1_log, T2_log)

    error_term = np.exp(T_max_log) * (np.exp(T1_log - T_max_log) + np.exp(T2_log - T_max_log))

    return error_term
//  [python]
def _extract_log_r(sim, i, n_thetas_marginal):
    log_p_xz_from_marginal = np.delete(sim.joint_log_probs, i, axis=0)
    delta_log = np.asarray(log_p_xz_from_marginal - sim.joint_log_probs[i] - np.log(float(n_thetas_marginal)), dtype=np.float128)
    log_r_xz = -1.0 * scipy.misc.logsumexp(delta_log)
    return log_r_xz
//  [python]
def _extract_log_r(sim, i, n_thetas_marginal):
    log_p_xz_from_marginal = np.delete(sim.joint_log_probs, i, axis=0)
    delta_log = np.asarray(log_p_xz_from_marginal - sim.joint_log_probs[i] - np.log(float(n_thetas_marginal)), dtype=np.float128)
    log_r_xz = -1.0 * scipy.special.logsumexp(delta_log)
    return log_r_xz
//  [python]
def _exponential_mixture_log_likelihood(params, t, t_min, t_max, t_step=None):
    """Calculate the log likelihood of an exponential mixture distribution.

    The full log likelihood for a single observation is given by:
        log(L) = log( sum_i( exp( log(component_i) ) ) )

    where log(component_i) is output from `_exponential_mixture_log_likelihood_components()`

    Parameters
    ----------
    params : array_like
        array of model parameters (amplitude and lifetime per component)
    t : array_like
        dwelltime observations in seconds
    t_min, t_max : float or np.ndarray
        minimum and maximum observation time in seconds
    t_step : float or np.ndarray, optional
        Discretization step
    """
    amplitudes, lifetimes = np.reshape(params, (2, -1))
    components = _exponential_mixture_log_likelihood_components(
        amplitudes, lifetimes, t, t_min, t_max, t_step
    )
    log_likelihood = np.log(np.sum(np.exp(components), axis=0))
    return -np.sum(log_likelihood)
//  [python]
def _exponential_mixture_log_likelihood(params, t, t_min, t_max, t_step=None):
    """Calculate the log likelihood of an exponential mixture distribution.

    The full log likelihood for a single observation is given by:
        log(L) = log( sum_i( exp( log(component_i) ) ) )

    where log(component_i) is output from `_exponential_mixture_log_likelihood_components()`

    Parameters
    ----------
    params : array_like
        array of model parameters (amplitude and lifetime per component)
    t : array_like
        dwelltime observations in seconds
    t_min, t_max : float or np.ndarray
        minimum and maximum observation time in seconds
    t_step : float or np.ndarray, optional
        Discretization step
    """
    amplitudes, lifetimes = np.reshape(params, (2, -1))
    components = _exponential_mixture_log_likelihood_components(
        amplitudes, lifetimes, t, t_min, t_max, t_step
    )
    log_likelihood = scipy.special.logsumexp(components, axis=0)
    return -np.sum(log_likelihood)
//  [python]
def calculate_pac_gamma(prob_np_all, label_np_all, alpha):
  """Compute the PAC-Gamma estimator."""
  starttime = time.time()
  label_np_all = one_hot(label_np_all)  # [n, v]
  soft_targets_sum = np.sum(label_np_all, axis=0)  # [v]
  soft_targets_sum = np.expand_dims(soft_targets_sum, axis=1)  # [v, 1]

  a0 = alpha * soft_targets_sum / np.sum(soft_targets_sum) + 1e-10
  beta = 1.

  # initialize
  qz = prob_np_all  # [n, d]
  s = prob_np_all  # [n, d]
  log_s = np.log(prob_np_all + 1e-10)  # [n, d]
  aw = a0
  bw = beta
  lw = np.sum(s, axis=-1, keepdims=True) * np.sum(aw / bw)  # [n, 1]

  for _ in range(10):
    aw = a0 + np.sum(np.einsum("BY,BZ->BYZ", label_np_all, qz),
                     axis=0)  # [v, d]
    lw = np.matmul(
        s, np.expand_dims(np.sum(aw / bw, axis=0), axis=1))  # [n, 1]
    logits_qz = (
        log_s + np.matmul(label_np_all, scipy.special.digamma(aw) - np.log(bw)))
    log_qz = logits_qz - np.max(logits_qz, axis=-1, keepdims=True)
    log_qz = log_qz - np.log(np.sum(np.exp(log_qz), axis=-1, keepdims=True))
    qz = np.exp(log_qz)  # [n, a, d]

  pac_gamma = (
      np.sum(scipy.special.loggamma(a0) - scipy.special.loggamma(aw) +
             aw * np.log(bw) - a0 * np.log(beta)) +
      np.sum(np.sum(qz * (log_qz - log_s), axis=-1) +
             np.log(np.squeeze(lw, axis=-1)) - 1.))
  pac_gamma /= label_np_all.size
  pac_gamma += 1.
  endtime = time.time()
  return [("pac_gamma_%.1f" % alpha, pac_gamma),
          ("time", endtime - starttime)]
//  [python]
def calculate_pac_gamma(prob_np_all, label_np_all, alpha):
  """Compute the PAC-Gamma estimator."""
  starttime = time.time()
  label_np_all = one_hot(label_np_all)  # [n, v]
  soft_targets_sum = np.sum(label_np_all, axis=0)  # [v]
  soft_targets_sum = np.expand_dims(soft_targets_sum, axis=1)  # [v, 1]

  a0 = alpha * soft_targets_sum / np.sum(soft_targets_sum) + 1e-10
  beta = 1.

  # initialize
  qz = prob_np_all  # [n, d]
  s = prob_np_all  # [n, d]
  log_s = np.log(prob_np_all + 1e-10)  # [n, d]
  aw = a0
  bw = beta
  lw = np.sum(s, axis=-1, keepdims=True) * np.sum(aw / bw)  # [n, 1]

  for _ in range(10):
    aw = a0 + np.sum(np.einsum("BY,BZ->BYZ", label_np_all, qz),
                     axis=0)  # [v, d]
    lw = np.matmul(
        s, np.expand_dims(np.sum(aw / bw, axis=0), axis=1))  # [n, 1]
    logits_qz = (
        log_s + np.matmul(label_np_all, scipy.special.digamma(aw) - np.log(bw)))
    log_qz = logits_qz - scipy.special.logsumexp(
        logits_qz, axis=-1, keepdims=True)
    qz = np.exp(log_qz)  # [n, a, d]

  pac_gamma = (
      np.sum(scipy.special.loggamma(a0) - scipy.special.loggamma(aw) +
             aw * np.log(bw) - a0 * np.log(beta)) +
      np.sum(np.sum(qz * (log_qz - log_s), axis=-1) +
             np.log(np.squeeze(lw, axis=-1)) - 1.))
  pac_gamma /= label_np_all.size
  pac_gamma += 1.
  endtime = time.time()
  return [("pac_gamma_%.1f" % alpha, pac_gamma),
          ("time", endtime - starttime)]
//  [python]
def act(self, obses, batch_size=32):
    def op(batch):
      feed_dict = {self.obs_ph: batch}
      return self.sess.run(self.eval_logits, feed_dict=feed_dict)
    logits = utils.batch_op(obses, batch_size, op)
    if self.discrete_act:
      logits -= scipy.misc.logsumexp(logits, axis=1, keepdims=True)
    return logits
//  [python]
def act(self, obses, batch_size=32):
    def op(batch):
      feed_dict = {self.obs_ph: batch}
      return self.sess.run(self.eval_logits, feed_dict=feed_dict)
    logits = utils.batch_op(obses, batch_size, op)
    if self.discrete_act:
      logits -= scipy.special.logsumexp(logits, axis=1, keepdims=True)
    return logits
//  [python]
def test_evidence(self):
        """Check evidence estimator."""
        self.assertAlmostEqual(e.evidence(self.ns_run),
                               np.exp(np.sum(np.log(self.w))),
                               places=12)
//  [python]
def test_evidence(self):
        """Check evidence estimator."""
        self.assertAlmostEqual(e.evidence(self.ns_run),
                               np.exp(scipy.special.logsumexp(self.logw)),
                               places=12)
//  [python]
def test_logx(self):
        """Check logx estimator."""
        self.assertAlmostEqual(e.logz(self.ns_run),
                               scipy.misc.logsumexp(self.logw), places=12)
//  [python]
def test_logx(self):
        """Check logx estimator."""
        self.assertAlmostEqual(e.logz(self.ns_run),
                               scipy.special.logsumexp(self.logw), places=12)
//  [python]
def testLogSumExpZeros(self):
    # Regression test for https://github.com/google/jax/issues/5370
    scipy_fun = lambda a, b: s.logsumexp(a, b=b)
    lax_fun = lambda a, b: lsp_special.logsumexp(a, b=b)
    args_maker = lambda: [np.array([-1000, -2]), np.array([1, 0])]
    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker)
    self._CompileAndCheck(lax_fun, args_maker)
//  [python]
def testLogSumExpZeros(self):
    # Regression test for https://github.com/google/jax/issues/5370
    scipy_fun = lambda a, b: osp_special.logsumexp(a, b=b)
    lax_fun = lambda a, b: lsp_special.logsumexp(a, b=b)
    args_maker = lambda: [np.array([-1000, -2]), np.array([1, 0])]
    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker)
    self._CompileAndCheck(lax_fun, args_maker)
//  [python]
def log_normalize(logprobs: dict[T, float]) -> dict[T, float]:
    """Normalize a log probability distribution so that all probabilities sum to 1.

    The input is a sparse distribution represented as a dictionary mapping from token IDs to log probabilities.
    """

    normalizer = cast(float, scipy.special.logsumexp(list(logprobs.values())))
    return {k: v - normalizer for k, v in logprobs.items()}
//  [python]
def _log_reconstruct_nz(self, subs_I_M):
        """Computes the log reconstruction for only non-zero entries."""
        I = subs_I_M[0].size
        K = self.n_components
        log_nz_recon_IK = np.zeros((I, K))
        for m in range(self.n_modes):
            log_nz_recon_IK += np.log(self.G_DK_M[m][subs_I_M[m], :])

        self.log_nz_recon_I = np.max(log_nz_recon_IK, axis=1) + np.log(np.sum(np.exp(log_nz_recon_IK - np.max(log_nz_recon_IK, axis=1, keepdims=True)), axis=1))
        return self.log_nz_recon_I
//  [python]
def _log_reconstruct_nz(self, subs_I_M):
        """Computes the log reconstruction for only non-zero entries."""
        I = subs_I_M[0].size
        K = self.n_components
        log_nz_recon_IK = np.zeros((I, K))
        for m in range(self.n_modes):
            log_nz_recon_IK += np.log(self.G_DK_M[m][subs_I_M[m], :])

        self.log_nz_recon_I = sp.logsumexp(log_nz_recon_IK, axis=1)
        return self.log_nz_recon_I
//  [python]
def run_classifier(features_np_all, label_np_all,
                   lda_factor,
                   features_np_eval, label_np_eval, is_vqa=False):
  """Compute the PAC_Gauss score with diagonal variance."""
  nclasses = label_np_all.max() + 1
  label_np_1hot = one_hot(label_np_all)  # [n, v]

  bs = features_np_all.shape[0]
  d = features_np_all.shape[-1]
  lda = lda_factor * bs

  def predict_ans(theta, features):
    theta = np.reshape(theta, [d + 1, nclasses])
    w = theta[:d, :]
    b = theta[d:, :]
    logits = np.matmul(features, w) + b
    predict = np.argmax(logits, axis=-1)
    return predict

  def compute_vqa_acc(label, predict):
    label_1hot = np.array([one_hot(t, nclasses) for t in label])
    targets = np.sum(label_1hot, axis=1)
    predictions = one_hot(predict, nclasses)
    vqa_acc = np.mean(
        np.minimum(np.sum(predictions * targets, axis=1) / 3.333, 1))
    return vqa_acc

  # optimizing log lik + log prior
  def pac_loss_fn(theta):
    theta = np.reshape(theta, [d + 1, nclasses])

    w = theta[:d, :]
    b = theta[d:, :]
    logits = np.matmul(features_np_all, w) + b

    log_qz = logits - np.max(logits, axis=-1, keepdims=True)
    log_qz -= np.log(np.sum(np.exp(log_qz), axis=-1, keepdims=True))
    xent = np.sum(np.sum(
        label_np_1hot * (np.log(label_np_1hot + 1e-10) - log_qz), axis=-1)) / bs
    loss = xent + 0.5 * np.sum(np.square(w)) / lda
    return loss

  # gradient of xent + l2
  def pac_grad_fn(theta):
    theta = np.reshape(theta, [d + 1, nclasses])

    w = theta[:d, :]
    b = theta[d:, :]
    logits = np.matmul(features_np_all, w) + b

    # grad_f = scipy.special.softmax(logits, axis=-1)  # [n, k]
    max_logits = logits.max(axis=-1, keepdims=True)
    grad_f = np.exp(logits - max_logits)
    grad_f /= np.sum(grad_f, axis=-1, keepdims=True)

    grad_f -= label_np_1hot
    grad_f /= bs
    grad_w = np.matmul(features_np_all.transpose(), grad_f)  # [d, k]
    grad_w += w / lda

    grad_b = np.sum(grad_f, axis=0, keepdims=True)  # [1, k]
    grad = np.ravel(np.concatenate([grad_w, grad_b], axis=0))
    return grad

  kernel_shape = [d, nclasses]
  theta = np.random.normal(size=kernel_shape) * 0.03
  theta_1d = np.ravel(np.concatenate(
      [theta, np.zeros([1, nclasses])], axis=0))

  theta_1d = scipy.optimize.minimize(
      pac_loss_fn, theta_1d, method="L-BFGS-B",
      jac=pac_grad_fn,
      options=dict(maxiter=50), tol=1e-6).x

  predict_train = predict_ans(theta_1d, features_np_all)
  train_acc = np.mean(label_np_all == predict_train)
  if is_vqa:
    predict_eval = predict_ans(theta_1d, features_np_eval)
    eval_acc = compute_vqa_acc(label_np_eval, predict_eval)
  else:
    predict_eval = predict_ans(theta_1d, features_np_eval)
    eval_acc = np.mean(label_np_eval == predict_eval)

  return train_acc, eval_acc
//  [python]
def run_classifier(features_np_all, label_np_all,
                   lda_factor,
                   features_np_eval, label_np_eval, is_vqa=False):
  """Compute the PAC_Gauss score with diagonal variance."""
  nclasses = label_np_all.max() + 1
  label_np_1hot = one_hot(label_np_all)  # [n, v]

  bs = features_np_all.shape[0]
  d = features_np_all.shape[-1]
  lda = lda_factor * bs

  def predict_ans(theta, features):
    theta = np.reshape(theta, [d + 1, nclasses])
    w = theta[:d, :]
    b = theta[d:, :]
    logits = np.matmul(features, w) + b
    predict = np.argmax(logits, axis=-1)
    return predict

  def compute_vqa_acc(label, predict):
    label_1hot = np.array([one_hot(t, nclasses) for t in label])
    targets = np.sum(label_1hot, axis=1)
    predictions = one_hot(predict, nclasses)
    vqa_acc = np.mean(
        np.minimum(np.sum(predictions * targets, axis=1) / 3.333, 1))
    return vqa_acc

  # optimizing log lik + log prior
  def pac_loss_fn(theta):
    theta = np.reshape(theta, [d + 1, nclasses])

    w = theta[:d, :]
    b = theta[d:, :]
    logits = np.matmul(features_np_all, w) + b

    log_qz = logits - scipy.special.logsumexp(logits, axis=-1, keepdims=True)
    xent = np.sum(np.sum(
        label_np_1hot * (np.log(label_np_1hot + 1e-10) - log_qz), axis=-1)) / bs
    loss = xent + 0.5 * np.sum(np.square(w)) / lda
    return loss

  # gradient of xent + l2
  def pac_grad_fn(theta):
    theta = np.reshape(theta, [d + 1, nclasses])

    w = theta[:d, :]
    b = theta[d:, :]
    logits = np.matmul(features_np_all, w) + b

    # grad_f = scipy.special.softmax(logits, axis=-1)  # [n, k]
    max_logits = logits.max(axis=-1, keepdims=True)
    grad_f = np.exp(logits - max_logits)
    grad_f /= np.sum(grad_f, axis=-1, keepdims=True)

    grad_f -= label_np_1hot
    grad_f /= bs
    grad_w = np.matmul(features_np_all.transpose(), grad_f)  # [d, k]
    grad_w += w / lda

    grad_b = np.sum(grad_f, axis=0, keepdims=True)  # [1, k]
    grad = np.ravel(np.concatenate([grad_w, grad_b], axis=0))
    return grad

  kernel_shape = [d, nclasses]
  theta = np.random.normal(size=kernel_shape) * 0.03
  theta_1d = np.ravel(np.concatenate(
      [theta, np.zeros([1, nclasses])], axis=0))

  theta_1d = scipy.optimize.minimize(
      pac_loss_fn, theta_1d, method="L-BFGS-B",
      jac=pac_grad_fn,
      options=dict(maxiter=50), tol=1e-6).x

  predict_train = predict_ans(theta_1d, features_np_all)
  train_acc = np.mean(label_np_all == predict_train)
  if is_vqa:
    predict_eval = predict_ans(theta_1d, features_np_eval)
    eval_acc = compute_vqa_acc(label_np_eval, predict_eval)
  else:
    predict_eval = predict_ans(theta_1d, features_np_eval)
    eval_acc = np.mean(label_np_eval == predict_eval)

  return train_acc, eval_acc
//  [python]
def testLogSumExp(self, shapes, dtype, axis,
                    keepdims, return_sign, use_b):
    if jtu.device_under_test() != "cpu":
      rng = jtu.rand_some_inf_and_nan(self.rng())
    else:
      rng = jtu.rand_default(self.rng())
    # TODO(mattjj): test autodiff
    if use_b:
      def scipy_fun(array_to_reduce, scale_array):
        res = s.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                    return_sign=return_sign, b=scale_array)
        if dtype == np.int32:
          res = tree_map(lambda x: x.astype('float32'), res)
        return res

      def lax_fun(array_to_reduce, scale_array):
        return lsp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                     return_sign=return_sign, b=scale_array)

      args_maker = lambda: [rng(shapes[0], dtype), rng(shapes[1], dtype)]
    else:
      def scipy_fun(array_to_reduce):
        res = s.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                    return_sign=return_sign)
        if dtype == np.int32:
          res = tree_map(lambda x: x.astype('float32'), res)
        return res

      def lax_fun(array_to_reduce):
        return lsp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                     return_sign=return_sign)

      args_maker = lambda: [rng(shapes[0], dtype)]
    tol = {np.float32: 1E-6, np.float64: 1E-14}
    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker)
    self._CompileAndCheck(lax_fun, args_maker, rtol=tol, atol=tol)
//  [python]
def testLogSumExp(self, shapes, dtype, axis,
                    keepdims, return_sign, use_b):
    if jtu.device_under_test() != "cpu":
      rng = jtu.rand_some_inf_and_nan(self.rng())
    else:
      rng = jtu.rand_default(self.rng())
    # TODO(mattjj): test autodiff
    if use_b:
      def scipy_fun(array_to_reduce, scale_array):
        res = osp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                    return_sign=return_sign, b=scale_array)
        if dtype == np.int32:
          res = tree_map(lambda x: x.astype('float32'), res)
        return res

      def lax_fun(array_to_reduce, scale_array):
        return lsp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                     return_sign=return_sign, b=scale_array)

      args_maker = lambda: [rng(shapes[0], dtype), rng(shapes[1], dtype)]
    else:
      def scipy_fun(array_to_reduce):
        res = osp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                    return_sign=return_sign)
        if dtype == np.int32:
          res = tree_map(lambda x: x.astype('float32'), res)
        return res

      def lax_fun(array_to_reduce):
        return lsp_special.logsumexp(array_to_reduce, axis, keepdims=keepdims,
                                     return_sign=return_sign)

      args_maker = lambda: [rng(shapes[0], dtype)]
    tol = {np.float32: 1E-6, np.float64: 1E-14}
    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker)
    self._CompileAndCheck(lax_fun, args_maker, rtol=tol, atol=tol)
//  [python]
def test_logsumexp(self, device):
        from scipy.special import logsumexp
        a = torch.randn(5, 4, device=device)
        a[0, 0] = inf
        a[1, :] = -inf
        actual = a.logsumexp(1)
        expected = logsumexp(a.cpu().numpy(), 1)
        self.assertEqual(expected.shape, actual.shape)
        self.assertEqual(expected, actual)

        # check that out is actually inplace
        b = torch.zeros(5, 2, device=device)
        c = b[:, 0]
        torch.logsumexp(a, 1, out=c)
        self.assertEqual(expected, b[:, 0])

        # check integral inputs is promoted to floating point
        e = torch.randint(-100, 100, [5, 4], device=device)
        actual = e.logsumexp(1).to(torch.float64)
        expected = logsumexp(e.cpu().numpy(), 1)
        self.assertEqual(expected.shape, actual.shape)
        self.assertEqual(expected, actual)
//  [python]
def test_logsumexp_dim(self, device):
        from scipy.special import logsumexp as s
        self._test_dim_ops(
            lambda t, d: t.logsumexp(d),
            lambda n, d: s(n, d),
            use_integral=False)
//  [python]
def test_logsumexp_dim(self, device):
        from scipy.special import logsumexp
        self._test_dim_ops(
            lambda t, d: t.logsumexp(d),
            lambda n, d: logsumexp(n, d),
            use_integral=False)
//  [python]
def remove_logit_bias(
    biased_logprobs: dict[T, float], logit_bias: dict[T, float] | dict[T, int]
) -> dict[T, float]:
    """Undo the effects of logit_bias.

    This function is useful to get log probabilities of arbitrary tokens,
    as a workaround for how OpenAI's API only returns the top K most likely tokens.
    Give those arbitary tokens a large logit_bias (about 70 works well in practice).

    For this function to work properly:
    - The logit_bias should be set so that logprobs for all tokens in the logit_bias
      are provided in biased_logprobs. For example, they should all be large positive numbers.
    - The logit_bias should not be too big, because of limited floating point precision.
      OpenAI's API will not return probabilities smaller than np.log(np.finfo(np.float32).smallest_normal.
    - biased_logprobs should have more elements than logit_bias. Otherwise it's
      not possible to recover the original log probabilities.
    """
    if not (logit_bias.keys() <= biased_logprobs.keys()):
        raise ValueError("logit_bias must be a subset of biased_logprobs")

    missing_key = object()
    missing_prob = 1 - np.exp(np.logaddexp.reduce(list(biased_logprobs.values())))
    missing_logprob = float(
        np.log(missing_prob) if missing_prob > MINIMUM_MISSING_PROB else -np.inf
    )

    debiased_logits: dict[T | object, float] = {
        k: v - logit_bias.get(k, 0) for k, v in biased_logprobs.items()
    }
    debiased_logits[missing_key] = missing_logprob

    normalized = log_normalize(debiased_logits)
    del normalized[missing_key]
    return cast(dict[T, float], normalized)
//  [python]
def remove_logit_bias(
    biased_logprobs: dict[T, float], logit_bias: dict[T, float] | dict[T, int]
) -> dict[T, float]:
    """Undo the effects of logit_bias.

    This function is useful to get log probabilities of arbitrary tokens,
    as a workaround for how OpenAI's API only returns the top K most likely tokens.
    Give those arbitary tokens a large logit_bias (about 70 works well in practice).

    For this function to work properly:
    - The logit_bias should be set so that logprobs for all tokens in the logit_bias
      are provided in biased_logprobs. For example, they should all be large positive numbers.
    - The logit_bias should not be too big, because of limited floating point precision.
      OpenAI's API will not return probabilities smaller than np.log(np.finfo(np.float32).smallest_normal.
    - biased_logprobs should have more elements than logit_bias. Otherwise it's
      not possible to recover the original log probabilities.
    """
    if not (logit_bias.keys() <= biased_logprobs.keys()):
        raise ValueError("logit_bias must be a subset of biased_logprobs")

    missing_key = object()
    missing_prob = 1 - np.exp(scipy.special.logsumexp(list(biased_logprobs.values())))
    missing_logprob = float(
        np.log(missing_prob) if missing_prob > MINIMUM_MISSING_PROB else -np.inf
    )

    debiased_logits: dict[T | object, float] = {
        k: v - logit_bias.get(k, 0) for k, v in biased_logprobs.items()
    }
    debiased_logits[missing_key] = missing_logprob

    normalized = log_normalize(debiased_logits)
    del normalized[missing_key]
    return cast(dict[T, float], normalized)
//  [python]
def _get_cluster_ccfs(self):
        n_samples = len(self.sample_list)
        cluster_ccfs = {}
        for mut in self.mutations.values():
            c = mut.cluster_assignment
            if c is not None:
                cluster_ccfs.setdefault(c, np.zeros((n_samples, 101)))
                cluster_ccfs[c] += np.log(mut.ccf_dist + 1e-10)
        for c in cluster_ccfs:
            cluster_ccfs[c] = np.exp(cluster_ccfs[c] - logsumexp(cluster_ccfs[c], axis=1, keepdims=True))
        return cluster_ccfs
//  [python]
def get_pi_dist(self):
        """
        Get pi dist by multiplying pi dists for supporting arm gains
        """
        pi_dist = np.zeros(101)
        for state in self.supporting_arm_states:
            for eve in state.cn_events:
                if eve.Type == 'Arm_gain':
                    if len(eve.supporting_muts) < self.min_timeable_muts:
                        continue
                    if eve.pi_dist is None:
                        eve.get_pi_dist_for_gain()
                    if eve.pi_dist is None:
                        continue
                    pi_dist += np.log(eve.pi_dist)
        self.pi_dist = np.exp(pi_dist - logsumexp(pi_dist))
//  [python]
def update(self):
        c_mllh = np.array([c.update() for c in self.children])
        lw = np.max(c_mllh) + np.log(np.sum(np.exp(c_mllh - np.max(c_mllh))))
        logweights = c_mllh - lw
        ## add a prior in case one child has too large a weight
        self.weights = np.exp(logweights) + np.ones(len(logweights)) * 0.5
        self.weights = self.weights / np.sum(self.weights)

        return lw
//  [python]
def update(self):
        c_mllh = np.array([c.update() for c in self.children])
        lw = logsumexp(c_mllh)
        logweights = c_mllh - lw
        ## add a prior in case one child has too large a weight
        self.weights = np.exp(logweights) + np.ones(len(logweights)) * 0.5
        self.weights = self.weights / np.sum(self.weights)

        return lw
//  [python]
def calcEofZCombII(EofLogZs, varLogZs, weights, space):
	"""
	calculate E(Z_t) from E(log(Z_i)) and Var(log(Z_i)) values where Z_t = sum_i W_i * Z_i
	space refers to output, not inputs. returning linear output runs the risk of under/overflow
	"""
	logEofZs = calcEofZII(EofLogZs, varLogZs, space = 'log')
	logWEofZs = np.log(weights) + logEofZs
	logWEofZ = np.logaddexp.reduce(logWEofZs)
	if space == 'linear':
		return np.exp(logWEofZ)
	elif space == 'log':
		return logWEofZ
//  [python]
def calcEofZCombII(EofLogZs, varLogZs, weights, space):
	"""
	calculate E(Z_t) from E(log(Z_i)) and Var(log(Z_i)) values where Z_t = sum_i W_i * Z_i
	space refers to output, not inputs. returning linear output runs the risk of under/overflow
	"""
	logEofZs = calcEofZII(EofLogZs, varLogZs, space = 'log')
	logWEofZs = np.log(weights) + logEofZs
	logWEofZ = scipy.special.logsumexp(logWEofZs)
	if space == 'linear':
		return np.exp(logWEofZ)
	elif space == 'log':
		return logWEofZ
//  [python]
def calcVarZCombII(EofLogZs, varLogZs, weights, space):
	logEofZs = calcEofZII(EofLogZs, varLogZs, space = 'log')
	logVarZs = calcVarZII(varLogZs, EofLogZs, logEofZs, space = 'log')
	logWVarZs = np.log(np.square(weights)) + logVarZs
	logWVarZ = np.max(logWVarZs) + np.log(np.sum(np.exp(logWVarZs - np.max(logWVarZs))))
	if space == 'linear':
		return np.exp(logWVarZ)
	elif space == 'log':
		return logWVarZ
//  [python]
def calcVarZCombII(EofLogZs, varLogZs, weights, space):
	logEofZs = calcEofZII(EofLogZs, varLogZs, space = 'log')
	logVarZs = calcVarZII(varLogZs, EofLogZs, logEofZs, space = 'log')
	logWVarZs = np.log(np.square(weights)) + logVarZs
	logWVarZ = scipy.special.logsumexp(logWVarZs) 
	if space == 'linear':
		return np.exp(logWVarZ)
	elif space == 'log':
		return logWVarZ
//  [python]
def get_p2_dist_for_gain(self):
        """
        get p2 distribution for a copy number gain
        """
        log_p2_posterior = np.zeros(101)
        for mut in self.supporting_muts:  # multiply mutation p2 distributions to get cnv p2 distribution
            mut_p2_dist = np.zeros(101)
            if (mut.local_cn_a1 == self.cn_a1).all() and (mut.local_cn_a2 == self.cn_a2).all() and mut.is_clonal:
                if mut.log_mult_dist is None:
                    mut.get_mult_dist()
                mut_p2_dist += np.sum(mut.log_mult_dist, 1)
            log_p2_posterior += mut_p2_dist - s.logsumexp(mut_p2_dist)
        log_p2_posterior = log_p2_posterior - s.logsumexp(log_p2_posterior) + self.log_p2_prior
        p2_real = self._correct_p2(log_p2_posterior - s.logsumexp(log_p2_posterior))
        return p2_real
//  [python]
def get_p2_dist_for_gain(self):
        """
        get p2 distribution for a copy number gain
        """
        log_p2_posterior = np.zeros(101)
        for mut in self.supporting_muts:  # multiply mutation p2 distributions to get cnv p2 distribution
            mut_p2_dist = np.zeros(101)
            if (mut.local_cn_a1 == self.cn_a1).all() and (mut.local_cn_a2 == self.cn_a2).all() and mut.is_clonal:
                if mut.log_mult_dist is None:
                    mut.get_mult_dist()
                mut_p2_dist += np.sum(mut.log_mult_dist, 1)
            log_p2_posterior += mut_p2_dist - logsumexp(mut_p2_dist)
        log_p2_posterior = log_p2_posterior - logsumexp(log_p2_posterior) + self.log_p2_prior
        p2_real = self._correct_p2(log_p2_posterior - logsumexp(log_p2_posterior))
        return p2_real
//  [python]
def pulsar_logweight(self, pulsar_masses, pulsar_masses_error, EOSpath, Neos):
        
        logweight = []
        for EOSIdx in range(1,Neos+1):
            EOSdata = '{0}/{1}.dat'.format(EOSpath, EOSIdx)
            R, M = np.loadtxt(EOSdata, usecols=[0,1], unpack=True)
            mMax = M[-1]

            logl = 0.
            for i in range(len(pulsar_masses)):
                mass = pulsar_masses[i]
                error = pulsar_masses_error[i]
                logl += scipy.stats.norm.logcdf(mMax, loc=mass, scale=error)
            logweight.append(logl)

        logweight -= scipy.misc.logsumexp(logweight)
        pulsar_weight = np.exp(logweight)
        np.savetxt('pulsar_weight.dat', np.c_[pulsar_weight])
        return pulsar_weight
//  [python]
def pulsar_logweight(self, pulsar_masses, pulsar_masses_error, EOSpath, Neos):
        
        logweight = []
        for EOSIdx in range(1,Neos+1):
            EOSdata = '{0}/{1}.dat'.format(EOSpath, EOSIdx)
            R, M = np.loadtxt(EOSdata, usecols=[0,1], unpack=True)
            mMax = M[-1]

            logl = 0.
            for i in range(len(pulsar_masses)):
                mass = pulsar_masses[i]
                error = pulsar_masses_error[i]
                logl += scipy.stats.norm.logcdf(mMax, loc=mass, scale=error)
            logweight.append(logl)

        logweight -= scipy.special.logsumexp(logweight)
        pulsar_weight = np.exp(logweight)
        np.savetxt('pulsar_weight.dat', np.c_[pulsar_weight])
        return pulsar_weight
//  [python]
def maxTOVmass_logweight(self, maxMTOV, maxMTOV_error, EOSpath, Neos):
        
        logweight = []
        for EOSIdx in range(1,Neos+1):
            EOSdata = '{0}/{1}.dat'.format(EOSpath, EOSIdx)
            R, M = np.loadtxt(EOSdata, usecols=[0,1], unpack=True)
            mMax = M[-1]

            logl = 0.
            for i in range(len(maxMTOV)):
                mass = maxMTOV[i]
                error = maxMTOV_error[i]
                logl += scipy.stats.norm.logsf(mMax, loc=mass, scale=error)
            logweight.append(logl)
        
        logweight -= scipy.misc.logsumexp(logweight)
        maxTOVmass_weight = np.exp(logweight)
        np.savetxt('maxTOVmass_weight.dat', np.c_[maxTOVmass_weight])

        return maxTOVmass_weight
//  [python]
def maxTOVmass_logweight(self, maxMTOV, maxMTOV_error, EOSpath, Neos):
        
        logweight = []
        for EOSIdx in range(1,Neos+1):
            EOSdata = '{0}/{1}.dat'.format(EOSpath, EOSIdx)
            R, M = np.loadtxt(EOSdata, usecols=[0,1], unpack=True)
            mMax = M[-1]

            logl = 0.
            for i in range(len(maxMTOV)):
                mass = maxMTOV[i]
                error = maxMTOV_error[i]
                logl += scipy.stats.norm.logsf(mMax, loc=mass, scale=error)
            logweight.append(logl)
        
        logweight -= scipy.special.logsumexp(logweight)
        maxTOVmass_weight = np.exp(logweight)
        np.savetxt('maxTOVmass_weight.dat', np.c_[maxTOVmass_weight])

        return maxTOVmass_weight
//  [python]
def log_likelihood(self, x, y):
        log_lik = self.log_complete_likelihood(x, y)
        return logsumexp(log_lik, axis=0)
//  [python]
def pulsar_mtov_weight(self, EOSpath, Neos, pulsar_weight, maxMTOV_weight):
        # eos weitghting with pulsar and mtov information, no NICER weighting
        total_weight = list(map(add, pulsar_weight, maxMTOV_weight))
        total_weight -= scipy.misc.logsumexp(total_weight)
        weight = np.exp(total_weight)
        np.savetxt('pulsar_mtov_weight.dat', np.c_[weight])
        
        return weight
//  [python]
def pulsar_mtov_weight(self, EOSpath, Neos, pulsar_weight, maxMTOV_weight):
        # eos weitghting with pulsar and mtov information, no NICER weighting
        total_weight = list(map(add, pulsar_weight, maxMTOV_weight))
        total_weight -= scipy.special.logsumexp(total_weight)
        weight = np.exp(total_weight)
        np.savetxt('pulsar_mtov_weight.dat', np.c_[weight])
        
        return weight
//  [python]
def total_logweight(self, EOSpath, Neos, pulsar_weight, maxMTOV_weight, NICER_weight):
        #eos weighting with all information: pulsar, mtov, NICER
        total_weight = list(map(add, pulsar_weight, maxMTOV_weight))
        total_weight = list(map(add, total_weight, NICER_weight))
        total_weight -= scipy.misc.logsumexp(total_weight)
        weight = np.exp(total_weight)
        np.savetxt('joint_weight.dat', np.c_[weight])
        
        return weight
//  [python]
def total_logweight(self, EOSpath, Neos, pulsar_weight, maxMTOV_weight, NICER_weight):
        #eos weighting with all information: pulsar, mtov, NICER
        total_weight = list(map(add, pulsar_weight, maxMTOV_weight))
        total_weight = list(map(add, total_weight, NICER_weight))
        total_weight -= scipy.special.logsumexp(total_weight)
        weight = np.exp(total_weight)
        np.savetxt('joint_weight.dat', np.c_[weight])
        
        return weight
//  [python]
def responsibilities(self, x, y):
        log_lik = self.log_complete_likelihood(x, y)
        resp = np.exp(log_lik - s.logsumexp(log_lik, axis=0, keepdims=True))
        return resp
//  [python]
def responsibilities(self, x, y):
        log_lik = self.log_complete_likelihood(x, y)
        resp = np.exp(log_lik - logsumexp(log_lik, axis=0, keepdims=True))
        return resp
//  [python]
def expected_log_likelihood(self, x, y):
        log_lik = self.expected_log_complete_likelihood(x, y)
        return logsumexp(log_lik, axis=0)