jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def logmeanexp(x, axis=None, **kwargs):
    """
    Evaluate the logarithm of the mean of exponentiated values in a numerically stable
    fashion.

    Parameters
    ----------
    x : array_like
        Input array.
    axis : None or int or tuple[int], optional
        Axis or axes over which the mean is taken. By default `axis` is None, and all
        elements are averaged.
    **kwargs : dict
        Keyword arguments passed to `scipy.special.logsumexp`.

    Returns
    -------
    y : ndarray
        `np.log(np.mean(np.exp(a)))` calculated in a numerically more stable way.
    """
    lse = special.logsumexp(x, axis, **kwargs)
    if isinstance(axis, collections.abc.Iterable):
        size = np.prod([np.size(x, a) for a in axis])
    else:
        size = np.size(x, axis)
    return lse - np.log(size)
//  [python]
def response_logprobabilities(self, query: Query) -> np.array:
        """Overwrites the parent's method. See :class:`.User` for more information."""
        if isinstance(query, PreferenceQuery):
            rewards = self.params['beta'] * self.reward(query.slate)
            return rewards - s.logsumexp(rewards)
            
        elif isinstance(query, WeakComparisonQuery):
            rewards = self.params['beta'] * self.reward(query.slate)
            logprobs = np.zeros((3))
            logprobs[1] = -np.log(1 + np.exp(self.params['delta'] + rewards[1] - rewards[0]))
            logprobs[2] = -np.log(1 + np.exp(self.params['delta'] + rewards[0] - rewards[1]))
            logprobs[0] = np.log(np.exp(2*self.params['delta']) - 1) + logprobs[1] + logprobs[2]
            return logprobs
            
        elif isinstance(query, FullRankingQuery):
            rewards = self.params['beta'] * self.reward(query.slate)
            logprobs = np.zeros(len(query.response_set))
            for response_id in range(len(query.response_set)):
                response = query.response_set[response_id]
                sorted_rewards = rewards[response]
                logprobs[response_id] = np.sum([sorted_rewards[i] - s.logsumexp(sorted_rewards[i:]) for i in range(len(response))])
            return logprobs
        raise NotImplementedError("response_logprobabilities is not defined for demonstration queries.")
//  [python]
def response_logprobabilities(self, query: Query) -> np.array:
        """Overwrites the parent's method. See :class:`.User` for more information."""
        if isinstance(query, PreferenceQuery):
            rewards = self.params['beta'] * self.reward(query.slate)
            return rewards - ssp.logsumexp(rewards)
            
        elif isinstance(query, WeakComparisonQuery):
            rewards = self.params['beta'] * self.reward(query.slate)
            logprobs = np.zeros((3))
            logprobs[1] = -np.log(1 + np.exp(self.params['delta'] + rewards[1] - rewards[0]))
            logprobs[2] = -np.log(1 + np.exp(self.params['delta'] + rewards[0] - rewards[1]))
            logprobs[0] = np.log(np.exp(2*self.params['delta']) - 1) + logprobs[1] + logprobs[2]
            return logprobs
            
        elif isinstance(query, FullRankingQuery):
            rewards = self.params['beta'] * self.reward(query.slate)
            logprobs = np.zeros(len(query.response_set))
            for response_id in range(len(query.response_set)):
                response = query.response_set[response_id]
                sorted_rewards = rewards[response]
                logprobs[response_id] = np.sum([sorted_rewards[i] - ssp.logsumexp(sorted_rewards[i:]) for i in range(len(response))])
            return logprobs
        raise NotImplementedError("response_logprobabilities is not defined for demonstration queries.")
//  [python]
def _ll(self, xs, cs, params):
    """Computes log likelihoods under the neural clustering process.
    """
    h_params, u_params, g_params, f_params = params
    num_data_points, unused_data_dim = xs.shape
    hs = vmap(jax_nn.nn_fwd, in_axes=(0, None))(xs, h_params)
    h_dim = hs.shape[1]
    us = vmap(jax_nn.nn_fwd, in_axes=(0, None))(xs, u_params)
    U = np.sum(us[1:], axis=0)
    Hs = np.zeros([num_data_points, h_dim])
    Hs = jax.ops.index_update(Hs, 0, hs[0])
    G = jax_nn.nn_fwd(Hs[0], g_params)
    K = 1

    def scan_body(state, i):
      U, K, G, Hs, ll = state
      U = U - us[i]

      def inner_map_body(k):
        H_k = Hs[k] + hs[i]
        G_k = G - jax_nn.nn_fwd(Hs[k], g_params) + jax_nn.nn_fwd(H_k, g_params)
        return self.f_fn(G_k, U, f_params)

      def map_body(k):
        return jax.lax.cond(np.less(k, K+1),
                            inner_map_body,
                            lambda x: -np.inf,
                            k)

      log_potentials = jax.lax.map(map_body, np.arange(num_data_points))
      log_Z_hat = s.logsumexp(log_potentials, keepdims=True)
      log_q = log_potentials - log_Z_hat
      ll += log_q[cs[i]]
      K = jax.lax.cond(np.equal(cs[i], K), lambda x: x + 1, lambda x: x, K)
      G = G - jax_nn.nn_fwd(Hs[cs[i]], g_params) + jax_nn.nn_fwd(
          Hs[cs[i]] + hs[i], g_params)
      Hs = jax.ops.index_update(Hs, cs[i], Hs[cs[i]] + hs[i])
      return (U, K, G, Hs, ll), None

    out = jax.lax.scan(
        scan_body,
        (U, K, G, Hs, 0.),
        np.arange(1, num_data_points))
    return out[0][4]
//  [python]
def _ll(self, xs, cs, params):
    """Computes log likelihoods under the neural clustering process.
    """
    h_params, u_params, g_params, f_params = params
    num_data_points, unused_data_dim = xs.shape
    hs = vmap(jax_nn.nn_fwd, in_axes=(0, None))(xs, h_params)
    h_dim = hs.shape[1]
    us = vmap(jax_nn.nn_fwd, in_axes=(0, None))(xs, u_params)
    U = np.sum(us[1:], axis=0)
    Hs = np.zeros([num_data_points, h_dim])
    Hs = jax.ops.index_update(Hs, 0, hs[0])
    G = jax_nn.nn_fwd(Hs[0], g_params)
    K = 1

    def scan_body(state, i):
      U, K, G, Hs, ll = state
      U = U - us[i]

      def inner_map_body(k):
        H_k = Hs[k] + hs[i]
        G_k = G - jax_nn.nn_fwd(Hs[k], g_params) + jax_nn.nn_fwd(H_k, g_params)
        return self.f_fn(G_k, U, f_params)

      def map_body(k):
        return jax.lax.cond(np.less(k, K+1),
                            inner_map_body,
                            lambda x: -np.inf,
                            k)

      log_potentials = jax.lax.map(map_body, np.arange(num_data_points))
      log_Z_hat = scipy.special.logsumexp(log_potentials, keepdims=True)
      log_q = log_potentials - log_Z_hat
      ll += log_q[cs[i]]
      K = jax.lax.cond(np.equal(cs[i], K), lambda x: x + 1, lambda x: x, K)
      G = G - jax_nn.nn_fwd(Hs[cs[i]], g_params) + jax_nn.nn_fwd(
          Hs[cs[i]] + hs[i], g_params)
      Hs = jax.ops.index_update(Hs, cs[i], Hs[cs[i]] + hs[i])
      return (U, K, G, Hs, ll), None

    out = jax.lax.scan(
        scan_body,
        (U, K, G, Hs, 0.),
        np.arange(1, num_data_points))
    return out[0][4]
//  [python]
def loglikelihood(self, data: QueryWithResponse) -> float:
        """
        Overwrites the parent's method. See :class:`.User` for more information.
        
        :Note: The loglikelihood value is the logarithm of the `unnormalized` likelihood if the
            input is a demonstration. Otherwise, it is the exact loglikelihood.
        """
        if isinstance(data, Demonstration):
            return self.params['beta_D'] * self.reward(data)
        
        elif isinstance(data, Preference):
            rewards = self.params['beta'] * self.reward(data.query.slate)
            return rewards[data.response] - s.logsumexp(rewards)
            
        elif isinstance(data, WeakComparison):
            rewards = self.params['beta'] * self.reward(data.query.slate)
            
            logp0 = -np.log(1 + np.exp(self.params['delta'] + rewards[1] - rewards[0]))
            if data.response == 0: return logp0
            
            logp1 = -np.log(1 + np.exp(self.params['delta'] + rewards[0] - rewards[1]))
            if data.response == 1: return logp1
            
            if data.response == -1:
                return np.log(np.exp(2*self.params['delta']) - 1) + logp0 + logp1
                
        elif isinstance(data, FullRanking):
            rewards = self.params['beta'] * self.reward(data.query.slate)
            sorted_rewards = rewards[data.response]
            return np.sum([sorted_rewards[i] - s.logsumexp(sorted_rewards[i:]) for i in range(len(data.response))])
            
        raise NotImplementedError("User response model for the given data is not implemented.")
//  [python]
def loglikelihood(self, data: QueryWithResponse) -> float:
        """
        Overwrites the parent's method. See :class:`.User` for more information.
        
        :Note: The loglikelihood value is the logarithm of the `unnormalized` likelihood if the
            input is a demonstration. Otherwise, it is the exact loglikelihood.
        """
        if isinstance(data, Demonstration):
            return self.params['beta_D'] * self.reward(data)
        
        elif isinstance(data, Preference):
            rewards = self.params['beta'] * self.reward(data.query.slate)
            return rewards[data.response] - ssp.logsumexp(rewards)
            
        elif isinstance(data, WeakComparison):
            rewards = self.params['beta'] * self.reward(data.query.slate)
            
            logp0 = -np.log(1 + np.exp(self.params['delta'] + rewards[1] - rewards[0]))
            if data.response == 0: return logp0
            
            logp1 = -np.log(1 + np.exp(self.params['delta'] + rewards[0] - rewards[1]))
            if data.response == 1: return logp1
            
            if data.response == -1:
                return np.log(np.exp(2*self.params['delta']) - 1) + logp0 + logp1
                
        elif isinstance(data, FullRanking):
            rewards = self.params['beta'] * self.reward(data.query.slate)
            sorted_rewards = rewards[data.response]
            return np.sum([sorted_rewards[i] - ssp.logsumexp(sorted_rewards[i:]) for i in range(len(data.response))])
            
        raise NotImplementedError("User response model for the given data is not implemented.")
//  [python]
def __call__(self, x, y, a, b):
    """Runs the Sinkhorn algorithm on two discrete measures (x, a) and (y, b).

    Args:
     x: np.ndarray<float>[batch, n] the 1D position of the input points.
     y: np.ndarray<float>[batch, m] the 1D position of the target points.
     a: np.ndarray<float>[batch, n] the weight of each input point.
     b: np.ndarray<float>[batch, m] the weight of each target point.

    Returns:
     np.ndarray<float>[batch, n, m]: the transport map between the two measures.
    """
    c = (x[:, :, np.newaxis] - y[:, np.newaxis, :]) ** self.power

    def center(f, g):
      return c - f[:, :, np.newaxis] - g[:, np.newaxis, :]

    def softmin(f, g, eps, axis):
      return -eps * np.logaddexp.reduce(-center(f, g) / eps, axis=axis)

    def error(f, g, eps):
      b_target = np.sum(np.exp(-center(f, g) / eps), axis=1)
      return np.max(np.abs(b_target - b) / b, axis=None)

    eps = self.epsilon_0
    loga = np.log(a)
    logb = np.log(b)
    f = np.zeros(np.shape(a), dtype=x.dtype)
    g = np.zeros(np.shape(b), dtype=x.dtype)
    err = self.threshold + 1.0
    self.iterations = 0
    while (self.iterations < self.max_iterations) and (
        err >= self.threshold or eps > self.epsilon):
      for _ in range(self.inner_iterations):
        self.iterations += 1
        g = eps * logb + softmin(f, g, eps, axis=1) + g
        f = eps * loga + softmin(f, g, eps, axis=2) + f
        eps = max(eps * self.epsilon_decay, self.epsilon)

      if eps <= self.epsilon:
        err = error(f, g, eps)

    return np.exp(-center(f, g) / eps)
//  [python]
def __call__(self, x, y, a, b):
    """Runs the Sinkhorn algorithm on two discrete measures (x, a) and (y, b).

    Args:
     x: np.ndarray<float>[batch, n] the 1D position of the input points.
     y: np.ndarray<float>[batch, m] the 1D position of the target points.
     a: np.ndarray<float>[batch, n] the weight of each input point.
     b: np.ndarray<float>[batch, m] the weight of each target point.

    Returns:
     np.ndarray<float>[batch, n, m]: the transport map between the two measures.
    """
    c = (x[:, :, np.newaxis] - y[:, np.newaxis, :]) ** self.power

    def center(f, g):
      return c - f[:, :, np.newaxis] - g[:, np.newaxis, :]

    def softmin(f, g, eps, axis):
      return -eps * scipy.special.logsumexp(-center(f, g) / eps, axis=axis)

    def error(f, g, eps):
      b_target = np.sum(np.exp(-center(f, g) / eps), axis=1)
      return np.max(np.abs(b_target - b) / b, axis=None)

    eps = self.epsilon_0
    loga = np.log(a)
    logb = np.log(b)
    f = np.zeros(np.shape(a), dtype=x.dtype)
    g = np.zeros(np.shape(b), dtype=x.dtype)
    err = self.threshold + 1.0
    self.iterations = 0
    while (self.iterations < self.max_iterations) and (
        err >= self.threshold or eps > self.epsilon):
      for _ in range(self.inner_iterations):
        self.iterations += 1
        g = eps * logb + softmin(f, g, eps, axis=1) + g
        f = eps * loga + softmin(f, g, eps, axis=2) + f
        eps = max(eps * self.epsilon_decay, self.epsilon)

      if eps <= self.epsilon:
        err = error(f, g, eps)

    return np.exp(-center(f, g) / eps)
//  [python]
def test_logsumexp(self):
        from scipy.misc import logsumexp as s
        a = torch.randn(5, 4)
        a[0, 0] = float('inf')
        a[1, :] = float('-inf')
        actual = a.logsumexp(1)
        expected = s(a.numpy(), 1)
        self.assertEqual(expected.shape, actual.shape)
        self.assertTrue(np.allclose(expected, actual.numpy()))
//  [python]
def test_logsumexp(self):
        from scipy.special import logsumexp
        a = torch.randn(5, 4)
        a[0, 0] = float('inf')
        a[1, :] = float('-inf')
        actual = a.logsumexp(1)
        expected = logsumexp(a.numpy(), 1)
        self.assertEqual(expected.shape, actual.shape)
        self.assertTrue(np.allclose(expected, actual.numpy()))
//  [python]
def _sample(self, xs, params, key):
    """Runs the neural clustering process.
    """
    h_params, u_params, g_params, f_params = params
    num_data_points, unused_data_dim = xs.shape
    hs = vmap(jax_nn.nn_fwd, in_axes=(0, None))(xs, h_params)
    h_dim = hs.shape[1]
    us = vmap(jax_nn.nn_fwd, in_axes=(0, None))(xs, u_params)
    U = np.sum(us[1:], axis=0)
    Hs = np.zeros([num_data_points, h_dim])
    Hs = jax.ops.index_update(Hs, 0, hs[0])
    G = jax_nn.nn_fwd(Hs[0], g_params)
    cs = np.zeros([num_data_points], dtype=np.int32)
    K = 1

    def for_body(i, state):
      U, K, G, Hs, cs, key = state
      U = U - us[i]

      def inner_map_body(k):
        H_k = Hs[k] + hs[i]
        G_k = G - jax_nn.nn_fwd(Hs[k], g_params) + jax_nn.nn_fwd(H_k, g_params)
        return self.f_fn(G_k, U, f_params)

      def map_body(k):
        return jax.lax.cond(np.less(k, K+1),
                            inner_map_body,
                            lambda x: -np.inf,
                            k)

      log_potentials = vmap(map_body)(np.arange(num_data_points))
      log_Z_hat = s.logsumexp(log_potentials, keepdims=True)
      log_q = log_potentials - log_Z_hat
      key, subkey = jax.random.split(key)
      c = jax.random.categorical(subkey, log_q)
      K = jax.lax.cond(np.equal(c, K), lambda x: x + 1, lambda x: x, K)
      G = G - jax_nn.nn_fwd(Hs[c], g_params) + jax_nn.nn_fwd(
          Hs[c] + hs[i], g_params)
      Hs = jax.ops.index_update(Hs, c, Hs[c] + hs[i])
      cs = jax.ops.index_update(cs, i, c)
      return (U, K, G, Hs, cs, key)

    out = jax.lax.fori_loop(1, num_data_points,
                            for_body,
                            (U, K, G, Hs, cs, jax.random.PRNGKey(0)))
    return out[4]
//  [python]
def _sample(self, xs, params, key):
    """Runs the neural clustering process.
    """
    h_params, u_params, g_params, f_params = params
    num_data_points, unused_data_dim = xs.shape
    hs = vmap(jax_nn.nn_fwd, in_axes=(0, None))(xs, h_params)
    h_dim = hs.shape[1]
    us = vmap(jax_nn.nn_fwd, in_axes=(0, None))(xs, u_params)
    U = np.sum(us[1:], axis=0)
    Hs = np.zeros([num_data_points, h_dim])
    Hs = jax.ops.index_update(Hs, 0, hs[0])
    G = jax_nn.nn_fwd(Hs[0], g_params)
    cs = np.zeros([num_data_points], dtype=np.int32)
    K = 1

    def for_body(i, state):
      U, K, G, Hs, cs, key = state
      U = U - us[i]

      def inner_map_body(k):
        H_k = Hs[k] + hs[i]
        G_k = G - jax_nn.nn_fwd(Hs[k], g_params) + jax_nn.nn_fwd(H_k, g_params)
        return self.f_fn(G_k, U, f_params)

      def map_body(k):
        return jax.lax.cond(np.less(k, K+1),
                            inner_map_body,
                            lambda x: -np.inf,
                            k)

      log_potentials = vmap(map_body)(np.arange(num_data_points))
      log_Z_hat = scipy.special.logsumexp(log_potentials, keepdims=True)
      log_q = log_potentials - log_Z_hat
      key, subkey = jax.random.split(key)
      c = jax.random.categorical(subkey, log_q)
      K = jax.lax.cond(np.equal(c, K), lambda x: x + 1, lambda x: x, K)
      G = G - jax_nn.nn_fwd(Hs[c], g_params) + jax_nn.nn_fwd(
          Hs[c] + hs[i], g_params)
      Hs = jax.ops.index_update(Hs, c, Hs[c] + hs[i])
      cs = jax.ops.index_update(cs, i, c)
      return (U, K, G, Hs, cs, key)

    out = jax.lax.fori_loop(1, num_data_points,
                            for_body,
                            (U, K, G, Hs, cs, jax.random.PRNGKey(0)))
    return out[4]
//  [python]
def _logpmf2(observed_counts, counts, base_lambda, error_rate):
        # Identical to _logpmf2, some lower memory usage by combining stuff
        poisson_lambda = (np.arange(counts.shape[1])[None, :] + error_rate) * base_lambda
        poisson_lambda = poisson_lambda.astype(np.float16)

        prob = s.logsumexp(np.log(counts / np.sum(counts, axis=-1)[:, None]) +
                         fast_poisson_logpmf(observed_counts[:, None].astype(np.float16), poisson_lambda)
                         , axis=-1)
        # prob where there are no counts should be 0 (-inf)
        prob[np.sum(counts, axis=-1) == 0] = -np.inf
        return prob
//  [python]
def _logpmf2(observed_counts, counts, base_lambda, error_rate):
        # Identical to _logpmf2, some lower memory usage by combining stuff
        poisson_lambda = (np.arange(counts.shape[1])[None, :] + error_rate) * base_lambda
        poisson_lambda = poisson_lambda.astype(np.float16)

        prob = logsumexp(np.log(counts / np.sum(counts, axis=-1)[:, None]) +
                         fast_poisson_logpmf(observed_counts[:, None].astype(np.float16), poisson_lambda)
                         , axis=-1)
        # prob where there are no counts should be 0 (-inf)
        prob[np.sum(counts, axis=-1) == 0] = -np.inf
        return prob
//  [python]
def test_logsumexp(self):
        val1 = beer.logsumexp(self.data, dim=0).numpy()
        val2 = s.logsumexp(self.data.numpy(), axis=0)
        self.assertArraysAlmostEqual(val1, val2)
        val1 = beer.logsumexp(self.data, dim=1).numpy()
        val2 = s.logsumexp(self.data.numpy(), axis=1)
        self.assertArraysAlmostEqual(val1, val2)
//  [python]
def test_logsumexp(self):
        val1 = beer.logsumexp(self.data, dim=0).numpy()
        val2 = logsumexp(self.data.numpy(), axis=0)
        self.assertArraysAlmostEqual(val1, val2)
        val1 = beer.logsumexp(self.data, dim=1).numpy()
        val2 = logsumexp(self.data.numpy(), axis=1)
        self.assertArraysAlmostEqual(val1, val2)
//  [python]
def logpmf(self, observed_counts, d, error_rate=0.0):
        kmer_counts, diplo_counts = self.diplotype_counts[d]
        diplo_counts = diplo_counts.astype(float)
        kmer_counts = kmer_counts.astype(float) + error_rate
        s = diplo_counts.sum(axis=-1)
        p_diplo_count = np.log(diplo_counts/s[:, None])
        broadcasted_counts = kmer_counts.shape.broadcast_values(observed_counts[:, None])
        p_kmer_counts = RaggedArray(poisson.logpmf(broadcasted_counts, kmer_counts.ravel()), kmer_counts.shape)
        ps = p_kmer_counts+p_diplo_count
        return np.array([s(row) for row in ps])
//  [python]
def logpmf(self, observed_counts, d, error_rate=0.0):
        kmer_counts, diplo_counts = self.diplotype_counts[d]
        diplo_counts = diplo_counts.astype(float)
        kmer_counts = kmer_counts.astype(float) + error_rate
        s = diplo_counts.sum(axis=-1)
        p_diplo_count = np.log(diplo_counts/s[:, None])
        broadcasted_counts = kmer_counts.shape.broadcast_values(observed_counts[:, None])
        p_kmer_counts = RaggedArray(poisson.logpmf(broadcasted_counts, kmer_counts.ravel()), kmer_counts.shape)
        ps = p_kmer_counts+p_diplo_count
        return np.array([logsumexp(row) for row in ps])
//  [python]
def update_beliefs(env, belief_logits, obs, prev_act, obs_model, dynamics_model, eps=1e-9):
  if env.name in ['gridworldnav', 'habitatnav']:
    assert len(obs_model.shape) == len(belief_logits.shape) + 1
    if len(obs_model.shape) == 2:
      shaped_obs_model = obs_model[np.newaxis, :, :]
      shaped_belief_logits = belief_logits[np.newaxis, :]
    elif len(obs_model.shape) == 3: # ensemble
      shaped_obs_model = obs_model
      shaped_belief_logits = belief_logits
    else:
      raise ValueError
    updated_belief_logits = deepcopy(shaped_obs_model[:, obs, :])
    if prev_act is not None:
      updated_belief_logits += np.logaddexp.reduce(
        dynamics_model[np.newaxis, :, :, prev_act] + shaped_belief_logits[:, np.newaxis, :], axis=2)
    updated_belief_logits -= np.logaddexp.reduce(updated_belief_logits, axis=1, keepdims=True)
    if len(obs_model.shape) == 2: # no ensemble
      updated_belief_logits = updated_belief_logits[0, :]
    return np.maximum(updated_belief_logits, -1e9)
  else:
    raise NotImplementedError
//  [python]
def update_beliefs(env, belief_logits, obs, prev_act, obs_model, dynamics_model, eps=1e-9):
  if env.name in ['gridworldnav', 'habitatnav']:
    assert len(obs_model.shape) == len(belief_logits.shape) + 1
    if len(obs_model.shape) == 2:
      shaped_obs_model = obs_model[np.newaxis, :, :]
      shaped_belief_logits = belief_logits[np.newaxis, :]
    elif len(obs_model.shape) == 3: # ensemble
      shaped_obs_model = obs_model
      shaped_belief_logits = belief_logits
    else:
      raise ValueError
    updated_belief_logits = deepcopy(shaped_obs_model[:, obs, :])
    if prev_act is not None:
      updated_belief_logits += scipy.special.logsumexp(
        dynamics_model[np.newaxis, :, :, prev_act] + shaped_belief_logits[:, np.newaxis, :], axis=2)
    updated_belief_logits -= scipy.special.logsumexp(updated_belief_logits, axis=1, keepdims=True)
    if len(obs_model.shape) == 2: # no ensemble
      updated_belief_logits = updated_belief_logits[0, :]
    return np.maximum(updated_belief_logits, -1e9)
  else:
    raise NotImplementedError
//  [python]
def logsumexp(x):
    if _GPU_ENABLED:
        # NOTE This is a quick-and-dirty implementation
        # FIXME Should contribute to the cupy codebase
        xmax = cp.amax(x)
        t = cp.exp(x - xmax)
        return cp.asnumpy(cp.log(cp.sum(t)) + xmax)
    else:
        return s.logsumexp(x)
//  [python]
def logsumexp(x):
    if _GPU_ENABLED:
        # NOTE This is a quick-and-dirty implementation
        # FIXME Should contribute to the cupy codebase
        xmax = cp.amax(x)
        t = cp.exp(x - xmax)
        return cp.asnumpy(cp.log(cp.sum(t)) + xmax)
    else:
        return scipy.special.logsumexp(x)
//  [python]
def median_and_err(xs, alpha=0.05):
    """Compute the median with an error which construct a central
    `100 * (1 - alpha)`%-confidence region using the order statistics.

    Source:
        https://stats.stackexchange.com/a/160148

    Args:
        xs (vector): Data.
        alpha (scalar): Desired confidence level.

    Returns:
        scalar: Median.
        scalar: Error which constructs a central `100 * (1 - alpha)`%-confidence region.
    """
    xs = B.to_numpy(B.flatten(xs))
    xs = xs[~np.isnan(xs)]
    xs = np.sort(xs)
    n = len(xs)

    def compute_log_term(i):
        log_binom = sps.loggamma(n + 1) - sps.loggamma(i + 1)
        log_binom = log_binom - sps.loggamma(n - i + 1) - n * np.log(2)
        return log_binom

    lower, upper = n // 2, n // 2
    cum_sum = compute_log_term(n // 2)

    while cum_sum < np.log(1 - alpha):
        lower -= 1
        cum_sum = np.logaddexp(cum_sum, compute_log_term(lower))
        if cum_sum < np.log(1 - alpha):
            upper += 1
            cum_sum = np.logaddexp(cum_sum, compute_log_term(upper))

    median = np.median(xs)
    error = max(median - xs[lower], median - xs[upper])
    return np.median(xs), error
//  [python]
def median_and_err(xs, alpha=0.05):
    """Compute the median with an error which construct a central
    `100 * (1 - alpha)`%-confidence region using the order statistics.

    Source:
        https://stats.stackexchange.com/a/160148

    Args:
        xs (vector): Data.
        alpha (scalar): Desired confidence level.

    Returns:
        scalar: Median.
        scalar: Error which constructs a central `100 * (1 - alpha)`%-confidence region.
    """
    xs = B.to_numpy(B.flatten(xs))
    xs = xs[~np.isnan(xs)]
    xs = np.sort(xs)
    n = len(xs)

    def compute_log_term(i):
        log_binom = sps.loggamma(n + 1) - sps.loggamma(i + 1)
        log_binom = log_binom - sps.loggamma(n - i + 1) - n * np.log(2)
        return log_binom

    lower, upper = n // 2, n // 2
    cum_sum = compute_log_term(n // 2)

    while cum_sum < np.log(1 - alpha):
        lower -= 1
        cum_sum = sps.logsumexp([cum_sum, compute_log_term(lower)])
        if cum_sum < np.log(1 - alpha):
            upper += 1
            cum_sum = sps.logsumexp([cum_sum, compute_log_term(upper)])

    median = np.median(xs)
    error = max(median - xs[lower], median - xs[upper])
    return np.median(xs), error
//  [python]
def __compute_donsker_varadhan_representation_of_kl_div__(self, ft_p, ft_q, is_optimization_else_estimation=False, is_max_approx=False):

        if is_optimization_else_estimation:

            if is_max_approx:
                kl_div = torch.mean(ft_p) - torch.max(ft_q) + torch.log(torch.tensor(ft_q.shape[0]))
            else:
                kl_div = torch.mean(ft_p) - torch.logsumexp(ft_q, dim=0) + torch.log(torch.tensor(ft_q.shape[0]))
        else:
            assert isinstance(ft_p, (np.ndarray, np.generic))
            assert isinstance(ft_q, (np.ndarray, np.generic))

            if is_max_approx:
                kl_div = ft_p.mean() - np.max(ft_q) + math.log(ft_q.size)
            else:
                kl_div = ft_p.mean() - s.logsumexp(ft_q) + math.log(ft_q.size)

        return kl_div
//  [python]
def __compute_donsker_varadhan_representation_of_kl_div__(self, ft_p, ft_q, is_optimization_else_estimation=False, is_max_approx=False):

        if is_optimization_else_estimation:

            if is_max_approx:
                kl_div = torch.mean(ft_p) - torch.max(ft_q) + torch.log(torch.tensor(ft_q.shape[0]))
            else:
                kl_div = torch.mean(ft_p) - torch.logsumexp(ft_q, dim=0) + torch.log(torch.tensor(ft_q.shape[0]))
        else:
            assert isinstance(ft_p, (np.ndarray, np.generic))
            assert isinstance(ft_q, (np.ndarray, np.generic))

            if is_max_approx:
                kl_div = ft_p.mean() - np.max(ft_q) + math.log(ft_q.size)
            else:
                kl_div = ft_p.mean() - scipy_special.logsumexp(ft_q) + math.log(ft_q.size)

        return kl_div
//  [python]
def logpdf_GMM(X, gmm):
    # This function will compute the log-density of a GMM for a set of samples
    # contained in matrix X of shape (D, N), where D is the size of a sample
    # and N is the number of samples in X.
    # We define a matrix S with shape (M, N). Each row will contain the sub-class
    # conditional densities given component Gi for all samples xi
    S = np.zeros((len(gmm), X.shape[1]))
    for i in range(len(gmm)):
        # Compute log density
        S[i, :] = logpdf_GAU_ND(X, gmm[i][1], gmm[i][2])
        # Add log of the prior of the corresponding component
        S[i, :] += np.log(gmm[i][0])
    # Compute the log-marginal log fxi(xi). The result will be an array of shape
    # (N,) whose component i will contain the log-density for sample xi
    logdens = sp.logsumexp(S, axis=0)
    return (logdens, S)
//  [python]
def logpdf_GMM(X, gmm):
    # This function will compute the log-density of a GMM for a set of samples
    # contained in matrix X of shape (D, N), where D is the size of a sample
    # and N is the number of samples in X.
    # We define a matrix S with shape (M, N). Each row will contain the sub-class
    # conditional densities given component Gi for all samples xi
    S = np.zeros((len(gmm), X.shape[1]))
    for i in range(len(gmm)):
        # Compute log density
        S[i, :] = logpdf_GAU_ND(X, gmm[i][1], gmm[i][2])
        # Add log of the prior of the corresponding component
        S[i, :] += np.log(gmm[i][0])
    # Compute the log-marginal log fxi(xi). The result will be an array of shape
    # (N,) whose component i will contain the log-density for sample xi
    logdens = scsp.logsumexp(S, axis=0)
    return (logdens, S)
//  [python]
def _normalize_logits(logits):
  numerator = logits
  denominator = s.logsumexp(logits)
  return numerator - denominator
//  [python]
def _normalize_logits(logits):
  numerator = logits
  denominator = scipy.special.logsumexp(logits)
  return numerator - denominator
//  [python]
def call(self, *args, **kwargs):
        return s(*args, **kwargs)
//  [python]
def call(self, *args, **kwargs):
        return scipy.special.logsumexp(*args, **kwargs)
//  [python]
def predict_proba(self, X):
        """
        Compute calibrated posterior probabilities for a given array of posterior probabilities from an arbitrary
        classifier.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_classes)
            The uncalibrated posterior probabilities.

        Returns
        -------
        P : array, shape (n_samples, n_classes)
            The predicted probabilities.
        """
        if X.ndim == 1:
            raise ValueError("Calibration data must have shape (n_samples, n_classes).")
        elif np.shape(X)[1] == 2:
            check_is_fitted(self, ["binnings", "log_scores", "prob_class_1", "T"])

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

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

            # Compute probability for other class
            return np.column_stack([1 - posterior_prob, posterior_prob])
        elif np.shape(X)[1] > 2:
            check_is_fitted(self, "onevsrest_calibrator_")
            return self.onevsrest_calibrator_.predict_proba(X)
//  [python]
def predict_proba(self, X):
        """
        Compute calibrated posterior probabilities for a given array of posterior probabilities from an arbitrary
        classifier.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_classes)
            The uncalibrated posterior probabilities.

        Returns
        -------
        P : array, shape (n_samples, n_classes)
            The predicted probabilities.
        """
        if X.ndim == 1:
            raise ValueError("Calibration data must have shape (n_samples, n_classes).")
        elif np.shape(X)[1] == 2:
            check_is_fitted(self, ["binnings", "log_scores", "prob_class_1", "T"])

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

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

            # Compute probability for other class
            return np.column_stack([1 - posterior_prob, posterior_prob])
        elif np.shape(X)[1] > 2:
            check_is_fitted(self, "onevsrest_calibrator_")
            return self.onevsrest_calibrator_.predict_proba(X)
//  [python]
def _log_mean(logx):
    # compute log of mean of x from log(x)
    return s.logsumexp(logx, axis=0) - np.log(len(logx))
//  [python]
def _log_mean(logx):
    # compute log of mean of x from log(x)
    return special.logsumexp(logx, axis=0) - np.log(len(logx))
//  [python]
def LSE_Sinkhorn(C, reg, a, b, num_iter=1000, delta=1e-3, lam=0):

    f = np.zeros(np.shape(a)[0])
    g = np.zeros(np.shape(b)[0])

    n, m = np.shape(C)

    C_tilde = f[:, None] + g[None, :] - C  # 2 * n * m
    C_tilde = C_tilde / reg  # n * m
    P = np.exp(C_tilde)

    err = 1
    n_iter = 0
    while n_iter < num_iter:
        P_prev = P
        if err > delta:
            n_iter = n_iter + 1

            # Update f
            f = reg * np.log(a) + f - reg * np.max(C_tilde, axis=1) - reg * np.log(np.sum(np.exp(C_tilde - np.max(C_tilde, axis=1)[:,None]), axis=1))

            # Update g
            C_tilde = f[:, None] + g[None, :] - C
            C_tilde = C_tilde / reg
            g = reg * np.log(b) + g - reg * np.max(C_tilde, axis=0) - reg * np.log(np.sum(np.exp(C_tilde - np.max(C_tilde, axis=0)[None,:]), axis=0))

            if (
                np.any(np.isnan(f))
                or np.any(np.isnan(g))
                or np.any(np.isinf(f))
                or np.any(np.isinf(g))
            ):
                # we have reached the machine precision
                # come back to previous solution and quit loop
                print("Warning: numerical errors at iteration", n_iter)
                P = P_prev
                break

            # Update the error
            C_tilde = f[:, None] + g[None, :] - C
            C_tilde = C_tilde / reg
            P = np.exp(C_tilde)
            err = np.sum(np.abs(np.sum(P, axis=1) - a))

        else:
            num_op = 4 * n * m + (n_iter + 1) * (8 * n * m + 3 * n + 3 * m) + n * m
            return P, num_op

    num_op = 4 * n * m + (n_iter + 1) * (8 * n * m + 3 * n + 3 * m) + n * m
    return P, num_op
//  [python]
def LSE_Sinkhorn(C, reg, a, b, num_iter=1000, delta=1e-3, lam=0):

    f = np.zeros(np.shape(a)[0])
    g = np.zeros(np.shape(b)[0])

    n, m = np.shape(C)

    C_tilde = f[:, None] + g[None, :] - C  # 2 * n * m
    C_tilde = C_tilde / reg  # n * m
    P = np.exp(C_tilde)

    err = 1
    n_iter = 0
    while n_iter < num_iter:
        P_prev = P
        if err > delta:
            n_iter = n_iter + 1

            # Update f
            f = reg * np.log(a) + f - reg * scipy.special.logsumexp(C_tilde, axis=1)

            # Update g
            C_tilde = f[:, None] + g[None, :] - C
            C_tilde = C_tilde / reg
            g = reg * np.log(b) + g - reg * scipy.special.logsumexp(C_tilde, axis=0)

            if (
                np.any(np.isnan(f))
                or np.any(np.isnan(g))
                or np.any(np.isinf(f))
                or np.any(np.isinf(g))
            ):
                # we have reached the machine precision
                # come back to previous solution and quit loop
                print("Warning: numerical errors at iteration", n_iter)
                P = P_prev
                break

            # Update the error
            C_tilde = f[:, None] + g[None, :] - C
            C_tilde = C_tilde / reg
            P = np.exp(C_tilde)
            err = np.sum(np.abs(np.sum(P, axis=1) - a))

        else:
            num_op = 4 * n * m + (n_iter + 1) * (8 * n * m + 3 * n + 3 * m) + n * m
            return P, num_op

    num_op = 4 * n * m + (n_iter + 1) * (8 * n * m + 3 * n + 3 * m) + n * m
    return P, num_op
//  [python]
def _log_var(logx):
    # compute log of variance of x from log(x)
    logmean = _log_mean(logx)
    pij = np.full_like(logx, np.pi * 1j, dtype=np.complex128)
    logxmu = s.logsumexp([logx, logmean + pij], axis=0)
    return np.real(s.logsumexp(2 * logxmu, axis=0)) - np.log(len(logx))
//  [python]
def _log_var(logx):
    # compute log of variance of x from log(x)
    logmean = _log_mean(logx)
    pij = np.full_like(logx, np.pi * 1j, dtype=np.complex128)
    logxmu = special.logsumexp([logx, logmean + pij], axis=0)
    return np.real(special.logsumexp(2 * logxmu, axis=0)) - np.log(len(logx))
//  [python]
def log_partition_ratio(
    x: np.ndarray,
    Phi_D_src: np.ndarray,
    n_src: int,
):
    """
    Calculate the log-partition ratio in the KLIEP problem.
    """
    return np.log(n_src) - scipy.special.logsumexp(Phi_D_src.dot(x))
//  [python]
def cross_ent_stable(y, logodds):
    '''
    issues with numerical stability of cross entropy

    if ytrue ==1 and p==0. logp = -inf and we have a problem 
    if we clip p[1e-10] = 1e-10 this is also a problem b/c the size of the clipping will
        drastically change the cross-entropy calculations

    solution: 
    pass in the logodds ratio log(P(A)/ (1-P(A))
    this logodds ratio is also the same as what comes out of sklearn logistic regression decision_function(X)
        
    Parameters
    ------
    D = number of instances 

    y: (D, 1) numpy array, true values of y 

    logodds: (D, ) numpy array of the log odds or the output of the LogisticRegression.decision_function
        which is equal to WX + b

    returns TOTAL entropy (you will have to manually divide to get average later on...)
    '''

    #if True in np.isinf(logp): raise ValueError('the logprobs have an inf or -inf value!')

    crossent = 0 
    D = y.shape[0]
    #assert D == logp.shape[0]
    assert D == logodds.shape[0]
    for i in xrange(D):
        #logodds = logp[i][1] - logp[i][0] #log(p/1-p) where p is pos class
        if y[i] == 1: 
            crossent += -1.0*np.logaddexp(0, -1.0*logodds[i])
        elif y[i] == 0 :
            crossent += -1.0*np.logaddexp(0, logodds[i])
        else: 
            raise ValueError('Your y array does not have exclusively 1s and 0s!')

    return -1.0*crossent
//  [python]
def cross_ent_stable(y, logodds):
    '''
    issues with numerical stability of cross entropy

    if ytrue ==1 and p==0. logp = -inf and we have a problem 
    if we clip p[1e-10] = 1e-10 this is also a problem b/c the size of the clipping will
        drastically change the cross-entropy calculations

    solution: 
    pass in the logodds ratio log(P(A)/ (1-P(A))
    this logodds ratio is also the same as what comes out of sklearn logistic regression decision_function(X)
        
    Parameters
    ------
    D = number of instances 

    y: (D, 1) numpy array, true values of y 

    logodds: (D, ) numpy array of the log odds or the output of the LogisticRegression.decision_function
        which is equal to WX + b

    returns TOTAL entropy (you will have to manually divide to get average later on...)
    '''

    #if True in np.isinf(logp): raise ValueError('the logprobs have an inf or -inf value!')

    crossent = 0 
    D = y.shape[0]
    #assert D == logp.shape[0]
    assert D == logodds.shape[0]
    for i in xrange(D):
        #logodds = logp[i][1] - logp[i][0] #log(p/1-p) where p is pos class
        if y[i] == 1: 
            crossent += -1.0*scipy.special.logsumexp([0, -1.0*logodds[i]])
        elif y[i] == 0 :
            crossent += -1.0*scipy.special.logsumexp([0, logodds[i]])
        else: 
            raise ValueError('Your y array does not have exclusively 1s and 0s!')

    return -1.0*crossent
//  [python]
def preference_based_on_gumble(fvals, indifference_threshold=0.0):
    # :param fvals: 1d-array of size num_choice
    assert len(fvals.shape) == 1

    num_choice = fvals.shape[0]

    indifference_mat = (1.0 - np.eye(num_choice)) * indifference_threshold

    base_f = fvals.reshape(-1,1) + indifference_mat * indifference_threshold

    pref_prob = np.exp(fvals - s.logsumexp(base_f, axis=-2))
    indifference_prob = 1.0 - np.sum(pref_prob)
    
    prob = np.concatenate([pref_prob, [indifference_prob]], axis=0)
    cprob = np.cumsum(prob)

    t = np.random.rand()
    selected_idx = np.searchsorted(cprob, t)

    if selected_idx == num_choice:
        # indifferent
        selected_idx = -1

    return selected_idx
//  [python]
def preference_based_on_gumble(fvals, indifference_threshold=0.0):
    # :param fvals: 1d-array of size num_choice
    assert len(fvals.shape) == 1

    num_choice = fvals.shape[0]

    indifference_mat = (1.0 - np.eye(num_choice)) * indifference_threshold

    base_f = fvals.reshape(-1,1) + indifference_mat * indifference_threshold

    pref_prob = np.exp(fvals - spmc.logsumexp(base_f, axis=-2))
    indifference_prob = 1.0 - np.sum(pref_prob)
    
    prob = np.concatenate([pref_prob, [indifference_prob]], axis=0)
    cprob = np.cumsum(prob)

    t = np.random.rand()
    selected_idx = np.searchsorted(cprob, t)

    if selected_idx == num_choice:
        # indifferent
        selected_idx = -1

    return selected_idx
//  [python]
def softtopk_np(logits, k):
    batchsize = logits.shape[0]
    f = softtopk_forward_np(logits, k)
    b = softtopk_backward_np(logits, k)
    initial_f = -INF * np.ones((batchsize, 1, k + 1))
    initial_f[:, :, 0] = 0
    ff = np.concatenate([initial_f, f[:, :-1, :]], axis=1)
    lse0 = s.logsumexp(ff + b, axis=2)
    lse1 = s.logsumexp(ff[:, :, :-1] + b[:, :, 1:], axis=2) + logits
    return np.exp(lse1 - np.logaddexp(lse0, lse1))
//  [python]
def softtopk_np(logits, k):
    batchsize = logits.shape[0]
    f = softtopk_forward_np(logits, k)
    b = softtopk_backward_np(logits, k)
    initial_f = -INF * np.ones((batchsize, 1, k + 1))
    initial_f[:, :, 0] = 0
    ff = np.concatenate([initial_f, f[:, :-1, :]], axis=1)
    lse0 = spec.logsumexp(ff + b, axis=2)
    lse1 = spec.logsumexp(ff[:, :, :-1] + b[:, :, 1:], axis=2) + logits
    return np.exp(lse1 - np.logaddexp(lse0, lse1))
//  [python]
def Sinkhorn_LSE(C, reg, a, b, max_iter=1000, delta=1e-3, lam=0, time_out=200):
    start = time.time()
    acc = []
    times = []

    C = C / C.max()
    n, m = np.shape(a)[0], np.shape(b)[0]

    f = np.zeros(n)
    g = np.zeros(m)

    C_trans = -C / reg
    P = np.exp(C_trans)
    OT_trans = np.sum(P * C)
    acc.append(OT_trans)

    time_actual = time.time() - start
    times.append(time_actual)

    err = 1
    n_iter = 0
    while n_iter < max_iter and (time_actual < time_out):
        P_prev = P
        if err > delta:
            n_iter = n_iter + 1

            # Update f
            C_tilde = f[:, None] + g[None, :] - C
            C_tilde = C_tilde / reg
            f = reg * np.log(a) + f - reg * np.max(C_tilde, axis=1) - reg * np.log(np.sum(np.exp(C_tilde - np.max(C_tilde, axis=1)[:,None]), axis=1))

            # Update g
            C_tilde = f[:, None] + g[None, :] - C
            C_tilde = C_tilde / reg
            g = reg * np.log(b) + g - reg * np.max(C_tilde, axis=0) - reg * np.log(np.sum(np.exp(C_tilde - np.max(C_tilde, axis=0)[None,:]), axis=0))

            # Update the coupling
            C_tilde = f[:, None] + g[None, :] - C
            C_tilde = C_tilde / reg
            P = np.exp(C_tilde)

            # Update the total cost
            OT_trans = np.sum(P * C)
            if np.isnan(OT_trans) == True:
                print("Error Sinkhorn: ", n_iter)
                P = P_prev
                break

            # Update the error
            err = np.sum(np.abs(np.sum(P, axis=1) - a))
            if np.isnan(err) == True:
                print("Error Sinkhorn: ", n_iter)
                P = P_prev
                break

            acc.append(OT_trans)
            time_actual = time.time() - start
            times.append(time_actual)

        else:
            break

    num_ops = 3 * n * m + (n_iter + 1) * (2 * n * m + n + m)
    return acc[-1], np.array(acc), np.array(times), P, num_ops
//  [python]
def Sinkhorn_LSE(C, reg, a, b, max_iter=1000, delta=1e-3, lam=0, time_out=200):
    start = time.time()
    acc = []
    times = []

    C = C / C.max()
    n, m = np.shape(a)[0], np.shape(b)[0]

    f = np.zeros(n)
    g = np.zeros(m)

    C_trans = -C / reg
    P = np.exp(C_trans)
    OT_trans = np.sum(P * C)
    acc.append(OT_trans)

    time_actual = time.time() - start
    times.append(time_actual)

    err = 1
    n_iter = 0
    while n_iter < max_iter and (time_actual < time_out):
        P_prev = P
        if err > delta:
            n_iter = n_iter + 1

            # Update f
            C_tilde = f[:, None] + g[None, :] - C
            C_tilde = C_tilde / reg
            f = reg * np.log(a) + f - reg * scipy.special.logsumexp(C_tilde, axis=1)

            # Update g
            C_tilde = f[:, None] + g[None, :] - C
            C_tilde = C_tilde / reg
            g = reg * np.log(b) + g - reg * scipy.special.logsumexp(C_tilde, axis=0)

            # Update the coupling
            C_tilde = f[:, None] + g[None, :] - C
            C_tilde = C_tilde / reg
            P = np.exp(C_tilde)

            # Update the total cost
            OT_trans = np.sum(P * C)
            if np.isnan(OT_trans) == True:
                print("Error Sinkhorn: ", n_iter)
                P = P_prev
                break

            # Update the error
            err = np.sum(np.abs(np.sum(P, axis=1) - a))
            if np.isnan(err) == True:
                print("Error Sinkhorn: ", n_iter)
                P = P_prev
                break

            acc.append(OT_trans)
            time_actual = time.time() - start
            times.append(time_actual)

        else:
            break

    num_ops = 3 * n * m + (n_iter + 1) * (2 * n * m + n + m)
    return acc[-1], np.array(acc), np.array(times), P, num_ops
//  [python]
def __call__(self, theta):
        """
        Calculate loglikelihood(theta), as well as any derived parameters.

        N.B. this loglikelihood is not normalised.

        Parameters
        ----------
        theta: float or 1d numpy array

        Returns
        -------
        logl: float
            Loglikelihood
        phi: list of length nderived
            Any derived parameters.
        """
        thetas = []
        for pos in self.positions:
            thetas.append(copy.deepcopy(theta))
            thetas[-1][:2] -= pos
        logls = [(Gaussian(sigma=self.sigmas[i])(thetas[i])[0]
                  + np.log(self.weights[i])) for i in range(len(self.weights))]
        logl = np.logaddexp.reduce(logls)
        return logl, [0.0] * self.nderived
//  [python]
def __call__(self, theta):
        """
        Calculate loglikelihood(theta), as well as any derived parameters.

        N.B. this loglikelihood is not normalised.

        Parameters
        ----------
        theta: float or 1d numpy array

        Returns
        -------
        logl: float
            Loglikelihood
        phi: list of length nderived
            Any derived parameters.
        """
        thetas = []
        for pos in self.positions:
            thetas.append(copy.deepcopy(theta))
            thetas[-1][:2] -= pos
        logls = [(Gaussian(sigma=self.sigmas[i])(thetas[i])[0]
                  + np.log(self.weights[i])) for i in range(len(self.weights))]
        logl = scipy.special.logsumexp(logls)
        return logl, [0.0] * self.nderived