<commit_before>import numpy as np
import itertools as itt

from scipy.special import beta, gamma, eval_jacobi, factorial


class MultivariateJacobiOPE:
    """
    Multivariate Jacobi Orthogonal Polynomial Ensemble used by :cite:`BaHa16` for Monte Carlo with Determinantal Point Processes.

    Reference weight
    .. math::

        w(x) = \prod_{i=1}^{d} (1-x)^{a_i} (1+x)^{b_i}

    :param N:
        Number of points :py:attr:`N`:math:`\geq 2`
    :type N:
        int, (default 10)

    :param jacobi_params:
        Jacobi parameters associated to
    :type N:
        array_like

    .. seealso::

        - :cite:`Kon05`
        - :cite:`BaHa16`
    """

    def __init__(self, N, jacobi_params):

        self.N, self.jacobi_params = self._check_params(N, jacobi_params)

        self.d = jacobi_params.shape[0]  # dimension

        self.max_layer, self.ordering = compute_ordering(self.N, self.d)

        self.square_norms =\
            compute_square_norms(self.jacobi_params, self.ordering)

        self.rejection_bound =\
            compute_rejection_bound(self.jacobi_params, self.ordering)

    def _check_params(self, N, jacobi_params):
        """ Check that:

        - The number of points :math:`N \geq 2`
        - Jacobi parameters :math:`(a_i, b_i)_{i=1}^d \in [-0.5, 0.5]^d`.
        """
        if type(N) is not int or N < 2:
            return TypeError('Number of points N={} is not an integer or < 2'.format(N))

        if jacobi_params.ndim < 2:
            err = ('`jacobi_params.ndim = {}` < 2'.format(jacobi_params.ndim),
                'This class implements d-dimensional (d>1) Jacobi ensemble.',
                'For a 3-dimensional example please provide',
                '`jacobi_params`=array([[a_1, b_1], [a_2, b_2], [a_3, b_3]])',
                'For d=1, use the tridiagonal model, cf JacobiEnsemble!')
            return ValueError('\n'.join(err))

        if (-0.5 <= jacobi_params).all() and (jacobi_params <= 0.5).all():
            return N, jacobi_params
        else:
            raise ValueError('Jacobi parameters not in [-0.5, 0.5]^d, we have no guaranty')

    def K(self, X, Y=None):
        '''
        .. math::
            K(x, y) = \sum_{\alpha}
                        \frac{P_{\alpha}(x)P_{\alpha}(y)}
                             {\|P_{\alpha}\|^2}

        for :math:`P_{\alpha}(x) = \prod_{i=1}^d P_{\alpha_i}^{a_i, b_i}(x_i)`
        '''

        if X.ndim == 1:
            if Y is None:
                return np.sum(
                            np.prod(
                                eval_jacobi(self.ordering,
                                            self.jacobi_params[:, 0],
                                            self.jacobi_params[:, 1],
                                            X),
                                axis=1)**2\
                            / self.square_norms)

            elif Y.ndim == 1:
                return np.sum(
                            np.prod(
                                eval_jacobi(self.ordering,
                                            self.jacobi_params[:, 0],
                                            self.jacobi_params[:, 1],
                                            [X[None, :], Y[None, :]]),
                                axis=(0, 2))\
                            / self.square_norms)

            else:  # [K(X, Y_1), ..., K(X, Y_n)]
                pol_evals = np.prod(
                                eval_jacobi(self.ordering,
                                            self.jacobi_params[:, 0],
                                            self.jacobi_params[:, 1],
                                            np.vstack((Y, X))[:, None]),
                                axis=2)

                return np.sum(pol_evals[-1] * pol_evals[:-1]\
                                / self.square_norms, axis=1)

        else:
            print('Weird call of this function X.shape={}, Y.shape={}'.format(X.shape, Y.shape))

    def sample_from_proposal(self, a=0.5, b=0.5):

        return 2.0 * np.random.beta(a, b, size=self.d) - 1.0

    def eval_w_over_mu_eq(self, x):

        a_b = 0.5 + self.jacobi_params

        return np.pi**self.d * np.prod((1 - x)**a_b[:, 0] * (1 + x)**a_b[:, 1])

    def sample(self, nb_trials_max=2000):
        """ Rejection sampling with proposal :math:`\mu_{\text{eq}}^{\otimes d}` where<commit_msg>reshape .K() kernel method for more flexibility<commit_after>import numpy as np
import itertools as itt

from scipy.special import beta, gamma, eval_jacobi, factorial


class MultivariateJacobiOPE:
    """
    Multivariate Jacobi Orthogonal Polynomial Ensemble used by :cite:`BaHa16` for Monte Carlo with Determinantal Point Processes.

    Reference weight
    .. math::

        w(x) = \prod_{i=1}^{d} (1-x)^{a_i} (1+x)^{b_i}

    :param N:
        Number of points :py:attr:`N`:math:`\geq 2`
    :type N:
        int, (default 10)

    :param jacobi_params:
        Jacobi parameters associated to
    :type N:
        array_like

    .. seealso::

        - :cite:`Kon05`
        - :cite:`BaHa16`
    """

    def __init__(self, N, jacobi_params):

        self.N, self.jacobi_params = self._check_params(N, jacobi_params)

        self.d = jacobi_params.shape[0]  # dimension

        self.max_layer, self.ordering = compute_ordering(self.N, self.d)

        self.square_norms =\
            compute_square_norms(self.jacobi_params, self.ordering)

        self.rejection_bound =\
            compute_rejection_bound(self.jacobi_params, self.ordering)

    def _check_params(self, N, jacobi_params):
        """ Check that:

        - The number of points :math:`N \geq 2`
        - Jacobi parameters :math:`(a_i, b_i)_{i=1}^d \in [-0.5, 0.5]^d`.
        """
        if type(N) is not int or N < 2:
            return TypeError('Number of points N={} is not an integer or < 2'.format(N))

        if jacobi_params.ndim < 2:
            err = ('`jacobi_params.ndim = {}` < 2'.format(jacobi_params.ndim),
                'This class implements d-dimensional (d>1) Jacobi ensemble.',
                'For a 3-dimensional example please provide',
                '`jacobi_params`=array([[a_1, b_1], [a_2, b_2], [a_3, b_3]])',
                'For d=1, use the tridiagonal model, cf JacobiEnsemble!')
            return ValueError('\n'.join(err))

        if (-0.5 <= jacobi_params).all() and (jacobi_params <= 0.5).all():
            return N, jacobi_params
        else:
            raise ValueError('Jacobi parameters not in [-0.5, 0.5]^d, we have no guaranty')

    def K(self, X, Y=None):
        '''
        K(x) == K(x, x)
        K(x, y) == K(y, x)
        K(x, Y) == K(Y, x) = [K(x, y) for y in Y]
        K(X, Y) == [K(x, y) for x, y in zip(X, Y)]
        .. math::
            K(x, y) = \sum_{\alpha}
                        \frac{P_{\alpha}(x)P_{\alpha}(y)}
                             {\|P_{\alpha}\|^2}

        for :math:`P_{\alpha}(x) = \prod_{i=1}^d P_{\alpha_i}^{a_i, b_i}(x_i)`
        '''
        if Y is None:

            vector = X.ndim == 1

            return np.sum(
                        np.prod(
                                eval_jacobi(self.ordering,
                                            self.jacobi_params[:, 0],
                                            self.jacobi_params[:, 1],
                                            X if vector else X[:, None]),
                                axis=1 if vector else 2)**2\
                            / self.square_norms,
                        axis=0 if vector else 1)
        else:

            pol_evals = np.prod(
                                eval_jacobi(self.ordering,
                                            self.jacobi_params[:, 0],
                                            self.jacobi_params[:, 1],
                                            np.vstack((Y, X))[:, None]),
                                axis=2)

            size = Y.shape[0] if Y.ndim > 1 else 1

            return np.sum(pol_evals[:size] * pol_evals[size:]\
                                / self.square_norms,
                        axis=1)


    def sample_from_proposal(self, a=0.5, b=0.5):

        return 2.0 * np.random.beta(a, b, size=self.d) - 1.0

    def eval_w_over_mu_eq(self, x):

        a_b = 0.5 + self.jacobi_params

        return np.pi**self.d * np.prod((1 - x)**a_b[:, 0] * (1 + x)**a_b[:, 1])

    def sample(self, nb_trials_max=2000):
        """ Rejection sampling with proposal :math:`\mu_{\text{eq}}^{\otimes d}` where
{'fraction_non_alphanumeric': 0.13571680583323534, 'fraction_numerical': 0.017876043749264964, 'mean_word_length': 1.4297142857142857, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 32, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4799862', 'n_tokens_mistral': 2831, 'n_tokens_neox': 2623, 'n_words': 608}