jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, nb_states, obs_dim, act_dim,
                 degree=1, **kwargs):

        self.nb_states = nb_states
        self.obs_dim = obs_dim
        self.act_dim = act_dim

        self.degree = degree
        self.feat_dim = int(comb(self.degree + self.obs_dim, self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        # self.K = npr.randn(self.nb_states, self.act_dim, self.feat_dim)
        # self.kff = npr.randn(self.nb_states, self.act_dim)
        # self._sigma_chol = 5. * npr.randn(self.nb_states, self.act_dim, self.act_dim)

        self.K = np.zeros((self.nb_states, self.act_dim, self.feat_dim))
        self.kff = np.zeros((self.nb_states, self.act_dim))
        self._sigma_chol = np.zeros((self.nb_states, self.act_dim, self.act_dim))

        for k in range(self.nb_states):
            _sigma = invwishart.rvs(self.act_dim + 1, np.eye(self.act_dim))
            self._sigma_chol[k] = np.linalg.cholesky(_sigma * np.eye(self.act_dim))
            self.K[k] = mvn.rvs(mean=None, cov=1e2 * _sigma, size=(self.feat_dim, )).T
            self.kff[k] = mvn.rvs(mean=None, cov=1e2 * _sigma, size=(1, ))
//  [python]
def __init__(self, nb_states, obs_dim, act_dim, nb_lags, prior, degree=1):

        assert nb_lags > 0

        self.nb_states = nb_states
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.nb_lags = nb_lags

        self.degree = degree
        self.feat_dim = int(s.comb(self.degree + (self.obs_dim * (self.nb_lags + 1)), self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        self.input_dim = self.feat_dim + 1
        self.output_dim = self.act_dim

        likelihood_precision_prior = prior['likelihood_precision_prior']
        parameter_precision_prior = prior['parameter_precision_prior']

        super().__init__(self.nb_states,
                         self.input_dim,
                         self.output_dim,
                         likelihood_precision_prior,
                         parameter_precision_prior)
//  [python]
def __init__(self, nb_states, obs_dim, act_dim, nb_lags, prior, degree=1):

        assert nb_lags > 0

        self.nb_states = nb_states
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.nb_lags = nb_lags

        self.degree = degree
        self.feat_dim = int(comb(self.degree + (self.obs_dim * (self.nb_lags + 1)), self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        self.input_dim = self.feat_dim + 1
        self.output_dim = self.act_dim

        likelihood_precision_prior = prior['likelihood_precision_prior']
        parameter_precision_prior = prior['parameter_precision_prior']

        super().__init__(self.nb_states,
                         self.input_dim,
                         self.output_dim,
                         likelihood_precision_prior,
                         parameter_precision_prior)
//  [python]
def __init__(self, nb_states, obs_dim, act_dim,
                 nb_lags, prior, degree=1, likelihood=None):

        assert nb_lags > 0

        self.nb_states = nb_states
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.nb_lags = nb_lags

        self.degree = degree
        self.feat_dim = int(s.comb(self.degree + (self.obs_dim * (self.nb_lags + 1)), self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        self.input_dim = self.feat_dim + 1
        self.output_dim = self.act_dim

        self.prior = prior
        self.posterior = copy.deepcopy(prior)

        # Linear-Gaussian likelihood
        if likelihood is not None:
            self.likelihood = likelihood
        else:
            As, lmbdas = self.prior.rvs()
            self.likelihood = StackedLinearGaussiansWithPrecision(size=self.nb_states,
                                                                  column_dim=self.input_dim,
                                                                  row_dim=self.output_dim,
                                                                  As=As, lmbdas=lmbdas, affine=True)
//  [python]
def __init__(self, nb_states, obs_dim, act_dim,
                 nb_lags, prior, degree=1, likelihood=None):

        assert nb_lags > 0

        self.nb_states = nb_states
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.nb_lags = nb_lags

        self.degree = degree
        self.feat_dim = int(comb(self.degree + (self.obs_dim * (self.nb_lags + 1)), self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        self.input_dim = self.feat_dim + 1
        self.output_dim = self.act_dim

        self.prior = prior
        self.posterior = copy.deepcopy(prior)

        # Linear-Gaussian likelihood
        if likelihood is not None:
            self.likelihood = likelihood
        else:
            As, lmbdas = self.prior.rvs()
            self.likelihood = StackedLinearGaussiansWithPrecision(size=self.nb_states,
                                                                  column_dim=self.input_dim,
                                                                  row_dim=self.output_dim,
                                                                  As=As, lmbdas=lmbdas, affine=True)
//  [python]
def _nCr(n, r):                                                                           # noqa
    """Number of combinations of r items out of a set of n.  Equals n!/(r!(n-r)!)"""
    f = _math.factorial; return f(n) / f(r) / f(n-r)
    #return _spspecial.comb(n, r)
//  [python]
def _nCr(n, r):                                                                           # noqa
    """Number of combinations of r items out of a set of n.  Equals n!/(r!(n-r)!)"""
    #f = _math.factorial; return f(n) / f(r) / f(n-r)
    return _spspecial.comb(n, r)
//  [python]
def __init__(self, nb_states, obs_dim, act_dim,
                 nb_lags=1, degree=1, **kwargs):

        assert nb_lags > 0

        self.nb_states = nb_states
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.nb_lags = nb_lags

        self.degree = degree
        self.feat_dim = int(s.comb(self.degree + (self.obs_dim * (self.nb_lags + 1)), self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        # self.K = npr.randn(self.nb_states, self.act_dim, self.feat_dim)
        # self.kff = npr.randn(self.nb_states, self.act_dim)
        # self._sigma_chol = 5. * npr.randn(self.nb_states, self.act_dim, self.act_dim)

        self.K = np.zeros((self.nb_states, self.act_dim, self.feat_dim))
        self.kff = np.zeros((self.nb_states, self.act_dim))
        self._sigma_chol = np.zeros((self.nb_states, self.act_dim, self.act_dim))

        for k in range(self.nb_states):
            _sigma = invwishart.rvs(self.act_dim + 1, np.eye(self.act_dim))
            self._sigma_chol[k] = np.linalg.cholesky(_sigma * np.eye(self.act_dim))
            self.K[k] = mvn.rvs(mean=None, cov=1e2 * _sigma, size=(self.feat_dim, )).T
            self.kff[k] = mvn.rvs(mean=None, cov=1e2 * _sigma, size=(1, ))
//  [python]
def __init__(self, nb_states, obs_dim, act_dim,
                 nb_lags=1, degree=1, **kwargs):

        assert nb_lags > 0

        self.nb_states = nb_states
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.nb_lags = nb_lags

        self.degree = degree
        self.feat_dim = int(comb(self.degree + (self.obs_dim * (self.nb_lags + 1)), self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        # self.K = npr.randn(self.nb_states, self.act_dim, self.feat_dim)
        # self.kff = npr.randn(self.nb_states, self.act_dim)
        # self._sigma_chol = 5. * npr.randn(self.nb_states, self.act_dim, self.act_dim)

        self.K = np.zeros((self.nb_states, self.act_dim, self.feat_dim))
        self.kff = np.zeros((self.nb_states, self.act_dim))
        self._sigma_chol = np.zeros((self.nb_states, self.act_dim, self.act_dim))

        for k in range(self.nb_states):
            _sigma = invwishart.rvs(self.act_dim + 1, np.eye(self.act_dim))
            self._sigma_chol[k] = np.linalg.cholesky(_sigma * np.eye(self.act_dim))
            self.K[k] = mvn.rvs(mean=None, cov=1e2 * _sigma, size=(self.feat_dim, )).T
            self.kff[k] = mvn.rvs(mean=None, cov=1e2 * _sigma, size=(1, ))
//  [python]
def bernstein_poly(n, i, t):
	"""
	Bernstein polynom.
	:param n: (int) polynom degree
	:param i: (int)
	:param t: (float)
	:return: (float)
	"""
	return s.comb(n, i) * t ** i * (1 - t) ** (n - i)
//  [python]
def bernstein_poly(n, i, t):
	"""
	Bernstein polynom.
	:param n: (int) polynom degree
	:param i: (int)
	:param t: (float)
	:return: (float)
	"""
	return scipy.special.comb(n, i) * t ** i * (1 - t) ** (n - i)
//  [python]
def betabinomial_pmf(k_list, size, a, b):
    """
        The beta binomial likelihood.
    """
    likelihood = 1

    for x in k_list:
        likelihood *= get_list_val(s.comb(size, x) * scipy.special.beta(
            x + a, size - x + b) / scipy.special.beta(a, b))

    return likelihood
//  [python]
def betabinomial_pmf(k_list, size, a, b):
    """
        The beta binomial likelihood.
    """
    likelihood = 1

    for x in k_list:
        likelihood *= get_list_val(scipy.special.comb(size, x) * scipy.special.beta(
            x + a, size - x + b) / scipy.special.beta(a, b))

    return likelihood
//  [python]
def __init__(self, nb_states, obs_dim, act_dim,
                 prior, degree, norm, device, **kwargs):
        super(SharedPolyOnlyRegressor, self).__init__(nb_states, obs_dim, act_dim,
                                                      prior, norm, device, **kwargs)

        self.degree = degree
        self.nb_feat = int(s.comb(self.degree + (self.obs_dim + self.act_dim), self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        self.weight = nn.Parameter(torch.Tensor(self.nb_states, self.nb_feat), requires_grad=True).to(self.device)
        self.bias = nn.Parameter(torch.Tensor(self.nb_states), requires_grad=True).to(self.device)

        self.reset_parameters()
//  [python]
def __init__(self, nb_states, obs_dim, act_dim,
                 prior, degree, norm, device, **kwargs):
        super(SharedPolyOnlyRegressor, self).__init__(nb_states, obs_dim, act_dim,
                                                      prior, norm, device, **kwargs)

        self.degree = degree
        self.nb_feat = int(sc.special.comb(self.degree + (self.obs_dim + self.act_dim), self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        self.weight = nn.Parameter(torch.Tensor(self.nb_states, self.nb_feat), requires_grad=True).to(self.device)
        self.bias = nn.Parameter(torch.Tensor(self.nb_states), requires_grad=True).to(self.device)

        self.reset_parameters()
//  [python]
def set_num_params(self, n_path=10000, dt=None, rn_seed=None, antithetic=True, m_inv=20, m_euler=20, n_euler=35):
        """
        Set MC parameters

        Args:
            n_path: number of paths
            dt: time step
            rn_seed: random number seed
            antithetic: antithetic
            m_inv: parameter M in Laplace inversion, Eq. (16)
            m_euler: parameter m in Euler transformation E(m,n)
            n_euler: parameter n in Euler transformation E(m,n)
        """
        self.m_inv = m_inv
        self.m_euler = m_euler
        self.n_euler = n_euler
        self.comb_coef = s.comb(self.m_euler, np.arange(0, self.m_euler+0.1)) * np.power(0.5, self.m_euler)
        assert abs(self.comb_coef.sum()-1) < 1e-8
        self.nn = np.arange(0, self.m_euler + self.n_euler + 0.1)
        super().set_num_params(n_path, dt, rn_seed, antithetic)
//  [python]
def set_num_params(self, n_path=10000, dt=None, rn_seed=None, antithetic=True, m_inv=20, m_euler=20, n_euler=35):
        """
        Set MC parameters

        Args:
            n_path: number of paths
            dt: time step
            rn_seed: random number seed
            antithetic: antithetic
            m_inv: parameter M in Laplace inversion, Eq. (16)
            m_euler: parameter m in Euler transformation E(m,n)
            n_euler: parameter n in Euler transformation E(m,n)
        """
        self.m_inv = m_inv
        self.m_euler = m_euler
        self.n_euler = n_euler
        self.comb_coef = spsp.comb(self.m_euler, np.arange(0, self.m_euler+0.1)) * np.power(0.5, self.m_euler)
        assert abs(self.comb_coef.sum()-1) < 1e-8
        self.nn = np.arange(0, self.m_euler + self.n_euler + 0.1)
        super().set_num_params(n_path, dt, rn_seed, antithetic)
//  [python]
def update_bezier_curve(self, i, j):
        self.add_control_point(i, j)

        if len(self.control_points) < 2:
            return

        def compute_bez_slice_points(control_points, num_points=252):
            n = len(control_points) - 1
            t = np.linspace(0, 1, num_points)
            slice_points = np.zeros((num_points, 2))
            for i in range(n + 1):
                binomial_coeff = s.comb(n, i)
                term = binomial_coeff * (t**i)[:, np.newaxis] * ((1 - t)**(n - i))[:, np.newaxis] * np.array(control_points[i])
                slice_points += term
            return slice_points

        self.slice_points = compute_bez_slice_points(self.control_points, num_points=252).T

        self.update_curve()
//  [python]
def update_bezier_curve(self, i, j):
        self.add_control_point(i, j)

        if len(self.control_points) < 2:
            return

        def compute_bez_slice_points(control_points, num_points=252):
            n = len(control_points) - 1
            t = np.linspace(0, 1, num_points)
            slice_points = np.zeros((num_points, 2))
            for i in range(n + 1):
                binomial_coeff = scipy.special.comb(n, i)
                term = binomial_coeff * (t**i)[:, np.newaxis] * ((1 - t)**(n - i))[:, np.newaxis] * np.array(control_points[i])
                slice_points += term
            return slice_points

        self.slice_points = compute_bez_slice_points(self.control_points, num_points=252).T

        self.update_curve()
//  [python]
def __init__(self,mean,var,d=50,p=2,device=None):
        #super(PCEModel, self).__init__()
        self.d = d
        self.p = p
        self.mean = mean
        self.var = var
        self.nbasis = s.comb(d+p,p).astype(int)
        if device==None:
            self.device = 'cpu'
        else:
            self.device = device
        self.oneDbasis = self.NormalBasis() 
        self.idxset = indexset(d,0)
        for i in range(1,p+1):
            self.idxset = torch.cat((self.idxset,indexset(d,i)),dim=0)
//  [python]
def __init__(self,mean,var,d=50,p=2,device=None):
        #super(PCEModel, self).__init__()
        self.d = d
        self.p = p
        self.mean = mean
        self.var = var
        self.nbasis = scipy.special.comb(d+p,p).astype(int)
        if device==None:
            self.device = 'cpu'
        else:
            self.device = device
        self.oneDbasis = self.NormalBasis() 
        self.idxset = indexset(d,0)
        for i in range(1,p+1):
            self.idxset = torch.cat((self.idxset,indexset(d,i)),dim=0)
//  [python]
def __init__(self, nb_states, obs_dim, act_dim, prior,
                 degree, norm, device, **kwargs):
        super(SharedPolyRegressor, self).__init__(nb_states, obs_dim, act_dim,
                                                  prior, norm, device, **kwargs)

        self.degree = degree

        self.nb_feat = int(s.comb(self.degree + (self.obs_dim + self.act_dim), self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        self.weight = nn.Parameter(torch.Tensor(self.nb_states, self.nb_feat), requires_grad=True).to(self.device)
        self.bias = nn.Parameter(torch.Tensor(self.nb_states), requires_grad=True).to(self.device)
        self.logmat = nn.Parameter(torch.Tensor(self.nb_states, self.nb_states), requires_grad=True).to(self.device)

        self.reset_parameters()
//  [python]
def __init__(self, nb_states, obs_dim, act_dim, prior,
                 degree, norm, device, **kwargs):
        super(SharedPolyRegressor, self).__init__(nb_states, obs_dim, act_dim,
                                                  prior, norm, device, **kwargs)

        self.degree = degree

        self.nb_feat = int(sc.special.comb(self.degree + (self.obs_dim + self.act_dim), self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        self.weight = nn.Parameter(torch.Tensor(self.nb_states, self.nb_feat), requires_grad=True).to(self.device)
        self.bias = nn.Parameter(torch.Tensor(self.nb_states), requires_grad=True).to(self.device)
        self.logmat = nn.Parameter(torch.Tensor(self.nb_states, self.nb_states), requires_grad=True).to(self.device)

        self.reset_parameters()
//  [python]
def bernstein(n, i, t):
    return s.comb(n, i) * t**i * (1 - t)**(n - i)
//  [python]
def bernstein(n, i, t):
    return scipy.special.comb(n, i) * t**i * (1 - t)**(n - i)
//  [python]
def _MatrixSufficientStatistics(
    x,
    y,
    epsilon: float,
    ) -> Tuple[int, int, int, int, int]:
    """Calculates tau sufficient statistics using matrices in NumPy.

    An absolute difference less than `epsilon` in x pairs is considered to be
    a tie.

    Args:
    x: Vector of numeric values.
    y: Vector of numeric values.
    epsilon: The threshold for which an absolute difference in x scores should
        be considered a tie.

    Returns:
    The number of concordant pairs, discordant pairs, pairs tied only in x,
    paired tied only in y, and pairs tied in both x and y.
    """
    x = np.asarray(x)
    x1, x2 = np.meshgrid(x, x.T)
    x_diffs = x1 - x2
    # Introduce ties into x by setting the diffs to 0 if they are <= epsilon
    x_is_tie = np.abs(x_diffs) <= epsilon
    x_diffs[x_is_tie] = 0.0
    
    y1, y2 = np.meshgrid(y, y.T)
    y_diffs = y1 - y2
    y_is_tie = y_diffs == 0.0

    n = len(y)
    num_pairs = int(s.comb(n, 2))
    # All of the counts are divided by 2 because each pair is double counted. The
    # double counted data will always be an even number, so dividing by 2 will
    # be an integer.
    con = int(
        ((x_diffs > 0) & (y_diffs > 0) | (x_diffs < 0) & (y_diffs < 0)).sum() / 2
    )
    t_x = int((x_is_tie & ~y_is_tie).sum() / 2)
    t_y = int((~x_is_tie & y_is_tie).sum() / 2)
    t_xy = int(((x_is_tie & y_is_tie).sum() - n) / 2)  # -n removes diagonal
    dis = num_pairs - (con + t_x + t_y + t_xy)
    return con, dis, t_x, t_y, t_xy
//  [python]
def _MatrixSufficientStatistics(
    x,
    y,
    epsilon: float,
    ) -> Tuple[int, int, int, int, int]:
    """Calculates tau sufficient statistics using matrices in NumPy.

    An absolute difference less than `epsilon` in x pairs is considered to be
    a tie.

    Args:
    x: Vector of numeric values.
    y: Vector of numeric values.
    epsilon: The threshold for which an absolute difference in x scores should
        be considered a tie.

    Returns:
    The number of concordant pairs, discordant pairs, pairs tied only in x,
    paired tied only in y, and pairs tied in both x and y.
    """
    x = np.asarray(x)
    x1, x2 = np.meshgrid(x, x.T)
    x_diffs = x1 - x2
    # Introduce ties into x by setting the diffs to 0 if they are <= epsilon
    x_is_tie = np.abs(x_diffs) <= epsilon
    x_diffs[x_is_tie] = 0.0
    
    y1, y2 = np.meshgrid(y, y.T)
    y_diffs = y1 - y2
    y_is_tie = y_diffs == 0.0

    n = len(y)
    num_pairs = int(scipy.special.comb(n, 2))
    # All of the counts are divided by 2 because each pair is double counted. The
    # double counted data will always be an even number, so dividing by 2 will
    # be an integer.
    con = int(
        ((x_diffs > 0) & (y_diffs > 0) | (x_diffs < 0) & (y_diffs < 0)).sum() / 2
    )
    t_x = int((x_is_tie & ~y_is_tie).sum() / 2)
    t_y = int((~x_is_tie & y_is_tie).sum() / 2)
    t_xy = int(((x_is_tie & y_is_tie).sum() - n) / 2)  # -n removes diagonal
    dis = num_pairs - (con + t_x + t_y + t_xy)
    return con, dis, t_x, t_y, t_xy
//  [python]
def __init__(
        self,
        inp_dim,
        n_vals=2,
        mix_strength=0,
        out_dim=None,
        total_power=None,
        n_units=None,
        mixing_order=None,
    ):
        if n_units is None:
            if mixing_order is None:
                order = inp_dim
            else:
                order = mixing_order
            n_units = int(inp_dim + s.comb(inp_dim, order) * n_vals**order)
        code = cc.make_code(
            1 - mix_strength,
            total_power,
            inp_dim,
            n_vals,
            n_units,
            mixing_order=mixing_order,
        )
        self.code = code
        self.generator = self.code.get_representation
        self.output_dim = n_units
        self.input_dim = inp_dim
//  [python]
def __init__(
        self,
        inp_dim,
        n_vals=2,
        mix_strength=0,
        out_dim=None,
        total_power=None,
        n_units=None,
        mixing_order=None,
    ):
        if n_units is None:
            if mixing_order is None:
                order = inp_dim
            else:
                order = mixing_order
            n_units = int(inp_dim + ss.comb(inp_dim, order) * n_vals**order)
        code = cc.make_code(
            1 - mix_strength,
            total_power,
            inp_dim,
            n_vals,
            n_units,
            mixing_order=mixing_order,
        )
        self.code = code
        self.generator = self.code.get_representation
        self.output_dim = n_units
        self.input_dim = inp_dim
//  [python]
def md_pattern(Y):
    """
    Get the missing data patterns from an array.

    Parameters
    ------------
    Y : :class:`numpy.ndarray` | :class:`pandas.DataFrame`
        Input dataset.

    Returns
    ---------
    pattern_ids : :class:`numpy.ndarray`
        Pattern ID array.
    pattern_dict : :class:`dict`
        Dictionary of patterns indexed by pattern IDs. Contains a pattern and count
        for each pattern ID.
    """
    if isinstance(Y, pd.DataFrame):
        Y = Y.values
    N, D = Y.shape
    # use int64 for higher-D arrays
    pID = np.zeros(N)
    Ymiss = ~np.isfinite(Y)
    rows = np.arange(N)[~np.isfinite(np.sum(Y, axis=1))]
    max_pats = s.comb(D, np.arange(0, D + 1)).sum()
    pID[rows] = max_pats * 5  # initialise to high value
    pD = defaultdict(dict)

    pindex = 0  # 0 = no missing data
    pD[int(0)] = {"pattern": np.zeros(D).astype(bool), "freq": np.sum(pID == 0)}
    indexes = np.arange(N)
    indexes = indexes[pID[indexes] > pindex]  # only look at md rows
    for idx in indexes:
        if pID[idx] > pindex:  # has missing data
            pindex += 1
            pID[idx] = pindex
            pattern = Ymiss[idx, :]
            pD[int(pindex)] = {"pattern": pattern, "freq": 0}
            if idx < N:
                _rix = np.arange(idx + 1, N)
                to_compare = _rix[pID[_rix] > pindex]
                where_same = to_compare[(Ymiss[to_compare, :] == pattern).all(axis=1)]
                pID[where_same] = pindex
    for ID in np.unique(pID):
        pD[ID]["freq"] = np.sum(pID == ID)
    return pID, pD
//  [python]
def md_pattern(Y):
    """
    Get the missing data patterns from an array.

    Parameters
    ------------
    Y : :class:`numpy.ndarray` | :class:`pandas.DataFrame`
        Input dataset.

    Returns
    ---------
    pattern_ids : :class:`numpy.ndarray`
        Pattern ID array.
    pattern_dict : :class:`dict`
        Dictionary of patterns indexed by pattern IDs. Contains a pattern and count
        for each pattern ID.
    """
    if isinstance(Y, pd.DataFrame):
        Y = Y.values
    N, D = Y.shape
    # use int64 for higher-D arrays
    pID = np.zeros(N)
    Ymiss = ~np.isfinite(Y)
    rows = np.arange(N)[~np.isfinite(np.sum(Y, axis=1))]
    max_pats = scipy.special.comb(D, np.arange(0, D + 1)).sum()
    pID[rows] = max_pats * 5  # initialise to high value
    pD = defaultdict(dict)

    pindex = 0  # 0 = no missing data
    pD[int(0)] = {"pattern": np.zeros(D).astype(bool), "freq": np.sum(pID == 0)}
    indexes = np.arange(N)
    indexes = indexes[pID[indexes] > pindex]  # only look at md rows
    for idx in indexes:
        if pID[idx] > pindex:  # has missing data
            pindex += 1
            pID[idx] = pindex
            pattern = Ymiss[idx, :]
            pD[int(pindex)] = {"pattern": pattern, "freq": 0}
            if idx < N:
                _rix = np.arange(idx + 1, N)
                to_compare = _rix[pID[_rix] > pindex]
                where_same = to_compare[(Ymiss[to_compare, :] == pattern).all(axis=1)]
                pID[where_same] = pindex
    for ID in np.unique(pID):
        pD[ID]["freq"] = np.sum(pID == ID)
    return pID, pD
//  [python]
def __init__(self, nb_states, obs_dim, act_dim,
                 nb_lags, prior, degree=1):

        assert nb_lags > 0

        self.nb_states = nb_states
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.nb_lags = nb_lags

        self.degree = degree
        self.feat_dim = int(s.comb(self.degree + self.obs_dim, self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        self.input_dim = self.feat_dim + 1
        self.output_dim = self.act_dim

        likelihood_precision_prior = prior['likelihood_precision_prior']
        parameter_precision_prior = prior['parameter_precision_prior']

        from sds.distributions.composite import StackedMultiOutputLinearGaussianWithAutomaticRelevance
        self.object = StackedMultiOutputLinearGaussianWithAutomaticRelevance(self.nb_states,
                                                                             self.input_dim,
                                                                             self.output_dim,
                                                                             likelihood_precision_prior,
                                                                             parameter_precision_prior)
//  [python]
def __init__(self, nb_states, obs_dim, act_dim,
                 nb_lags, prior, degree=1):

        assert nb_lags > 0

        self.nb_states = nb_states
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.nb_lags = nb_lags

        self.degree = degree
        self.feat_dim = int(sc.special.comb(self.degree + self.obs_dim, self.degree)) - 1
        self.basis = PolynomialFeatures(self.degree, include_bias=False)

        self.input_dim = self.feat_dim + 1
        self.output_dim = self.act_dim

        likelihood_precision_prior = prior['likelihood_precision_prior']
        parameter_precision_prior = prior['parameter_precision_prior']

        from sds.distributions.composite import StackedMultiOutputLinearGaussianWithAutomaticRelevance
        self.object = StackedMultiOutputLinearGaussianWithAutomaticRelevance(self.nb_states,
                                                                             self.input_dim,
                                                                             self.output_dim,
                                                                             likelihood_precision_prior,
                                                                             parameter_precision_prior)
//  [python]
def ER(n, p):
        """Generate an Erdos-Renyi graph

        Args:
            n (int): number of nodes
            p (float): probability of an edge

        Returns:
            nx.Graph: a graph
        """
        return transform._proba2graph(list(itertools.combinations(range(n),2)), [p]*s.comb(n,2,exact=True))
//  [python]
def ER(n, p):
        """Generate an Erdos-Renyi graph

        Args:
            n (int): number of nodes
            p (float): probability of an edge

        Returns:
            nx.Graph: a graph
        """
        return transform._proba2graph(list(itertools.combinations(range(n),2)), [p]*scipy.special.comb(n,2,exact=True))
//  [python]
def __init__(
        self,
        inp_dim,
        n_vals=2,
        mix_strength=0,
        out_dim=None,
        total_power=1,
        n_units=None,
        mixing_order=None,
        seed=None,
        n_stim_set=10000,
    ):
        self.input_dim = inp_dim
        self.rng = np.random.default_rng(seed)
        self.n_vals = n_vals
        self.total_power = total_power
        self.m_units_nz = s.comb(inp_dim, mixing_order) * n_vals**mixing_order
        categories = [np.arange(n_vals)] * inp_dim
        ohe = skp.OneHotEncoder(categories=categories, sparse_output=False)
        if mixing_order is None:
            mixing_order = inp_dim
        pf = skp.PolynomialFeatures(
            degree=(mixing_order, mixing_order),
            include_bias=False,
            interaction_only=True
        )
        scaler = skp.StandardScaler()
        make_mix_feats = sklpipe.make_pipeline(ohe, pf, scaler)
        fit_stim = self.sample_stim(n_stim_set)
        self.scaler = None
        if mix_strength > 0 and mixing_order > 1:
            self.mix_pipe = make_mix_feats.fit(fit_stim)
            m_len = self.mix_pipe.transform(fit_stim).shape[1]
        else:
            self.mix_pipe = None
            m_len = 0
        p_len = fit_stim.shape[1]
        self.pure_pipe = skp.StandardScaler().fit(fit_stim)
        if n_units is None:
            n_units = p_len + m_len
        self.output_dim = n_units
        self.p_units = p_len
        self.m_units = m_len
        self.trs_mat = sp.linalg.orth(
            self.rng.normal(size=(n_units, n_units))
        )[:, :p_len + m_len]
        self.trs_mat = u.make_unit_vector(self.trs_mat)
        self.mix_strength = mix_strength
//  [python]
def __init__(
        self,
        inp_dim,
        n_vals=2,
        mix_strength=0,
        out_dim=None,
        total_power=1,
        n_units=None,
        mixing_order=None,
        seed=None,
        n_stim_set=10000,
    ):
        self.input_dim = inp_dim
        self.rng = np.random.default_rng(seed)
        self.n_vals = n_vals
        self.total_power = total_power
        self.m_units_nz = ss.comb(inp_dim, mixing_order) * n_vals**mixing_order
        categories = [np.arange(n_vals)] * inp_dim
        ohe = skp.OneHotEncoder(categories=categories, sparse_output=False)
        if mixing_order is None:
            mixing_order = inp_dim
        pf = skp.PolynomialFeatures(
            degree=(mixing_order, mixing_order),
            include_bias=False,
            interaction_only=True
        )
        scaler = skp.StandardScaler()
        make_mix_feats = sklpipe.make_pipeline(ohe, pf, scaler)
        fit_stim = self.sample_stim(n_stim_set)
        self.scaler = None
        if mix_strength > 0 and mixing_order > 1:
            self.mix_pipe = make_mix_feats.fit(fit_stim)
            m_len = self.mix_pipe.transform(fit_stim).shape[1]
        else:
            self.mix_pipe = None
            m_len = 0
        p_len = fit_stim.shape[1]
        self.pure_pipe = skp.StandardScaler().fit(fit_stim)
        if n_units is None:
            n_units = p_len + m_len
        self.output_dim = n_units
        self.p_units = p_len
        self.m_units = m_len
        self.trs_mat = sp.linalg.orth(
            self.rng.normal(size=(n_units, n_units))
        )[:, :p_len + m_len]
        self.trs_mat = u.make_unit_vector(self.trs_mat)
        self.mix_strength = mix_strength
//  [python]
def _r_single(X, N=1000,limit=0.01):
    """
    Assume X is trials X time raster (channel removed)

    test data from SPN recording
    X=rec['resp'].extract_epoch('STIM_BNB+si464+si1889')[:, chanidx, :]
    """

    repcount = X.shape[0]
    if repcount <= 1:
        log.info('repcount<=1, rnorm=0')
        return 0

    paircount = np.int(scipy.misc.comb(repcount, 2))
    pairs = []
    for p1 in range(repcount):
        for p2 in range(p1+1, repcount):
            pairs.append([p1, p2])

    if paircount < N:
        N = paircount

    if N == 0:
        # TODO:
        # only two repeats, break up data in time to get a better
        # estimate of single-trial correlations
        # raise ValueError("2 repeats condition not supported yet.")
        # N=10;
        # bstep=size(pred,1)./N;
        # rac=zeros(N,1);
        # for nn=1:N,
        #     tt=round((nn-1).*bstep+1):round(nn*bstep);
        #     if ~isempty(tt) && std(resp(tt,1))>0 && std(resp(tt,2))>0,
        #         rac(nn)=xcov(resp(tt,1),resp(tt,2),0,'coeff');
        #     end
        # end
        # print('r_ceiling invalid')
        return 0.05
    else:

        rac = np.zeros(N)
        sidx = np.argsort(np.random.rand(paircount))
        for nn in range(N):
            X1 = X[pairs[sidx[nn]][0], :]
            X2 = X[pairs[sidx[nn]][1], :]

            # remove all nans from pred and resp
            ff = np.isfinite(X1) & np.isfinite(X2)
            X1 = X1[ff]
            X2 = X2[ff]
            #if (np.sum(X1) > 0) and (np.sum(X2) > 0):
            if    (np.sum(X1) == 0) or (np.sum(X2) == 0) or \
                   (len(np.unique(X1)) == 1) or \
                   (len(np.unique(X2)) == 1):
                rac[nn] = 0
            else:
                rac[nn] = np.corrcoef(X1, X2)[0, 1]
                

    # hard limit on single-trial correlation to prevent explosion
    # TODO: better logic for this
    rac = np.mean(rac)
    if rac < limit:
        rac = limit

    return rac
//  [python]
def _r_single(X, N=1000,limit=0.01):
    """
    Assume X is trials X time raster (channel removed)

    test data from SPN recording
    X=rec['resp'].extract_epoch('STIM_BNB+si464+si1889')[:, chanidx, :]
    """

    repcount = X.shape[0]
    if repcount <= 1:
        log.info('repcount<=1, rnorm=0')
        return 0

    paircount = np.int(scipy.special.comb(repcount, 2))
    pairs = []
    for p1 in range(repcount):
        for p2 in range(p1+1, repcount):
            pairs.append([p1, p2])

    if paircount < N:
        N = paircount

    if N == 0:
        # TODO:
        # only two repeats, break up data in time to get a better
        # estimate of single-trial correlations
        # raise ValueError("2 repeats condition not supported yet.")
        # N=10;
        # bstep=size(pred,1)./N;
        # rac=zeros(N,1);
        # for nn=1:N,
        #     tt=round((nn-1).*bstep+1):round(nn*bstep);
        #     if ~isempty(tt) && std(resp(tt,1))>0 && std(resp(tt,2))>0,
        #         rac(nn)=xcov(resp(tt,1),resp(tt,2),0,'coeff');
        #     end
        # end
        # print('r_ceiling invalid')
        return 0.05
    else:

        rac = np.zeros(N)
        sidx = np.argsort(np.random.rand(paircount))
        for nn in range(N):
            X1 = X[pairs[sidx[nn]][0], :]
            X2 = X[pairs[sidx[nn]][1], :]

            # remove all nans from pred and resp
            ff = np.isfinite(X1) & np.isfinite(X2)
            X1 = X1[ff]
            X2 = X2[ff]
            #if (np.sum(X1) > 0) and (np.sum(X2) > 0):
            if    (np.sum(X1) == 0) or (np.sum(X2) == 0) or \
                   (len(np.unique(X1)) == 1) or \
                   (len(np.unique(X2)) == 1):
                rac[nn] = 0
            else:
                rac[nn] = np.corrcoef(X1, X2)[0, 1]
                

    # hard limit on single-trial correlation to prevent explosion
    # TODO: better logic for this
    rac = np.mean(rac)
    if rac < limit:
        rac = limit

    return rac
//  [python]
def getUnlinkedStationary(self, popSize, theta):
        one_loc_probs = one_locus_probs(popSize=popSize, theta=theta, n=self.n)
        assertValidProbs(one_loc_probs)

        n = self.n
        leftOnes = self.numOnes(0)
        rightOnes = self.numOnes(1)
        bothOnes = self.hapCount((1, 1))
        joint = one_loc_probs[leftOnes] * one_loc_probs[rightOnes]
        if self.exact:
            joint[self.numC > 0] = 0
        else:
            joint *= (s.comb(rightOnes, bothOnes)
                      * s.comb(n - rightOnes, leftOnes - bothOnes)
                      / s.comb(n, leftOnes))

        joint *= self.n_unfolded_versions

        assertValidProbs(joint)
        return joint
//  [python]
def getUnlinkedStationary(self, popSize, theta):
        one_loc_probs = one_locus_probs(popSize=popSize, theta=theta, n=self.n)
        assertValidProbs(one_loc_probs)

        n = self.n
        leftOnes = self.numOnes(0)
        rightOnes = self.numOnes(1)
        bothOnes = self.hapCount((1, 1))
        joint = one_loc_probs[leftOnes] * one_loc_probs[rightOnes]
        if self.exact:
            joint[self.numC > 0] = 0
        else:
            joint *= (scipy.special.comb(rightOnes, bothOnes)
                      * scipy.special.comb(n - rightOnes, leftOnes - bothOnes)
                      / scipy.special.comb(n, leftOnes))

        joint *= self.n_unfolded_versions

        assertValidProbs(joint)
        return joint
//  [python]
def BinomialCoef(n, k):
    """Compute the binomial coefficient "n choose k".

    n: number of trials
    k: number of successes

    Returns: float
    """
    return s.comb(n, k)
//  [python]
def comb_index(n, k):
    count = s.comb(n, k, exact=True)
    index = np.fromiter(itertools.chain.from_iterable(itertools.combinations(range(n), k)), 
                        int, count=count*k)
    return index.reshape(-1, k)
//  [python]
def comb_index(n, k):
    count = scipy.special.comb(n, k, exact=True)
    index = np.fromiter(itertools.chain.from_iterable(itertools.combinations(range(n), k)), 
                        int, count=count*k)
    return index.reshape(-1, k)
//  [python]
def load_dataset_large(dataset_name):
    global number_spins
    magnetisation = number_spins % 2
    number_ups = (number_spins + magnetisation) // 2
    shift = number_ups * (number_ups - 1) // 2 if number_ups > 0 else 0
    all_spins = np.fromiter(
        sector(number_spins, magnetisation),
        dtype=np.uint64,
        count=int(scipy.misc.comb(number_spins, number_ups)),
    ).astype(np.int64)

    all_amplitudes = torch.from_numpy(load_eigenvector(dataset_name))
    print('norm = ', torch.sum(all_amplitudes ** 2).item())
    ipr = int(torch.sum(all_amplitudes ** 2).item() ** 2 / torch.sum(all_amplitudes ** 4).item())
    print('IPR = ', str(ipr))
    sort_ampls = np.sort(all_amplitudes)[-ipr:]
    print('1/IPR states contains weigth: ', np.sum(sort_ampls ** 2))
    # return -1
    dataset = (
        torch.from_numpy(all_spins),
        torch.where(all_amplitudes >= 0, torch.tensor([0]), torch.tensor([1])).squeeze(),
        (torch.abs(all_amplitudes) ** 2).type(torch.FloatTensor),
    )


    return dataset
//  [python]
def load_dataset_large(dataset_name):
    global number_spins
    magnetisation = number_spins % 2
    number_ups = (number_spins + magnetisation) // 2
    shift = number_ups * (number_ups - 1) // 2 if number_ups > 0 else 0
    all_spins = np.fromiter(
        sector(number_spins, magnetisation),
        dtype=np.uint64,
        count=int(scipy.special.comb(number_spins, number_ups)),
    ).astype(np.int64)

    all_amplitudes = torch.from_numpy(load_eigenvector(dataset_name))
    print('norm = ', torch.sum(all_amplitudes ** 2).item())
    ipr = int(torch.sum(all_amplitudes ** 2).item() ** 2 / torch.sum(all_amplitudes ** 4).item())
    print('IPR = ', str(ipr))
    sort_ampls = np.sort(all_amplitudes)[-ipr:]
    print('1/IPR states contains weigth: ', np.sum(sort_ampls ** 2))
    # return -1
    dataset = (
        torch.from_numpy(all_spins),
        torch.where(all_amplitudes >= 0, torch.tensor([0]), torch.tensor([1])).squeeze(),
        (torch.abs(all_amplitudes) ** 2).type(torch.FloatTensor),
    )


    return dataset
//  [python]
def discrete_gaussian_distribution(self, grid_height, grid_width):
        '''
        generate link_weight(rewards) with gaussian distribution
        :param grid_height: per grid height
        :param grid_width: per grid width
        :return:
        '''
        height = int(self.height / grid_height + 0.999999999999)
        width = int(self.width / grid_width + 0.999999999999)
        gd = np.zeros((height, width))
        for i in range(len(gd)):
            for j in range(len(gd[i])):
                gd[i, j] = s.comb(height - 1, i, exact=True) + s.comb(width - 1, j, exact=True)
        gd = np.float32(gd) / np.float32(gd).sum()
        return gd
//  [python]
def discrete_gaussian_distribution(self, grid_height, grid_width):
        '''
        generate link_weight(rewards) with gaussian distribution
        :param grid_height: per grid height
        :param grid_width: per grid width
        :return:
        '''
        height = int(self.height / grid_height + 0.999999999999)
        width = int(self.width / grid_width + 0.999999999999)
        gd = np.zeros((height, width))
        for i in range(len(gd)):
            for j in range(len(gd[i])):
                gd[i, j] = sps.comb(height - 1, i, exact=True) + sps.comb(width - 1, j, exact=True)
        gd = np.float32(gd) / np.float32(gd).sum()
        return gd
//  [python]
def __init__(self, mean, var, d=50, p=2, device=None):
        super(PCEModel, self).__init__()
        self.d = d
        self.p = p
        self.mean = mean
        self.var = var
        # scipy.special.comb(N,k):The number of combinations of N
        # things taken k at a time.
        self.nbasis = s.comb(d + p, p).astype(int)
        if device is None:
            self.device = 'cpu'
        else:
            self.device = device
        self.oneDbasis = self.NormalBasis()
        self.idxset = indexset(d, 0)
        for i in range(1, p + 1):
            self.idxset = torch.cat((self.idxset, indexset(d, i)), dim=0)

        self.mean = nn.Parameter(self.mean, requires_grad=False)
        self.var = nn.Parameter(self.var, requires_grad=False)
        self.oneDbasis = nn.Parameter(self.oneDbasis, requires_grad=False)
        self.idxset = nn.Parameter(self.idxset, requires_grad=False)
//  [python]
def __init__(self, mean, var, d=50, p=2, device=None):
        super(PCEModel, self).__init__()
        self.d = d
        self.p = p
        self.mean = mean
        self.var = var
        # scipy.special.comb(N,k):The number of combinations of N
        # things taken k at a time.
        self.nbasis = scipy.special.comb(d + p, p).astype(int)
        if device is None:
            self.device = 'cpu'
        else:
            self.device = device
        self.oneDbasis = self.NormalBasis()
        self.idxset = indexset(d, 0)
        for i in range(1, p + 1):
            self.idxset = torch.cat((self.idxset, indexset(d, i)), dim=0)

        self.mean = nn.Parameter(self.mean, requires_grad=False)
        self.var = nn.Parameter(self.var, requires_grad=False)
        self.oneDbasis = nn.Parameter(self.oneDbasis, requires_grad=False)
        self.idxset = nn.Parameter(self.idxset, requires_grad=False)
//  [python]
def _adjusted_rand_index(reference_indices, estimated_indices):
    """Compute the Rand index, adjusted for change.

    Parameters
    ----------
    reference_indices : np.ndarray
        Array of reference indices
    estimated_indices : np.ndarray
        Array of estimated indices

    Returns
    -------
    ari : float
        Adjusted Rand index
    .. note:: Based on sklearn.metrics.cluster.adjusted_rand_score

    """
    n_samples = len(reference_indices)
    ref_classes = np.unique(reference_indices)
    est_classes = np.unique(estimated_indices)
    # Special limit cases: no clustering since the data is not split;
    # or trivial clustering where each document is assigned a unique cluster.
    # These are perfect matches hence return 1.0.
    if (
        ref_classes.shape[0] == est_classes.shape[0] == 1
        or ref_classes.shape[0] == est_classes.shape[0] == 0
        or (ref_classes.shape[0] == est_classes.shape[0] == len(reference_indices))
    ):
        return 1.0

    contingency = _contingency_matrix(reference_indices, estimated_indices)

    # Compute the ARI using the contingency data
    sum_comb_c = sum(
        s.comb(n_c, 2, exact=1) for n_c in contingency.sum(axis=1)
    )
    sum_comb_k = sum(
        s.comb(n_k, 2, exact=1) for n_k in contingency.sum(axis=0)
    )

    sum_comb = sum(
        (s.comb(n_ij, 2, exact=1) for n_ij in contingency.flatten())
    )
    prod_comb = (sum_comb_c * sum_comb_k) / float(s.comb(n_samples, 2))
    mean_comb = (sum_comb_k + sum_comb_c) / 2.0
    return (sum_comb - prod_comb) / (mean_comb - prod_comb)
//  [python]
def _adjusted_rand_index(reference_indices, estimated_indices):
    """Compute the Rand index, adjusted for change.

    Parameters
    ----------
    reference_indices : np.ndarray
        Array of reference indices
    estimated_indices : np.ndarray
        Array of estimated indices

    Returns
    -------
    ari : float
        Adjusted Rand index
    .. note:: Based on sklearn.metrics.cluster.adjusted_rand_score

    """
    n_samples = len(reference_indices)
    ref_classes = np.unique(reference_indices)
    est_classes = np.unique(estimated_indices)
    # Special limit cases: no clustering since the data is not split;
    # or trivial clustering where each document is assigned a unique cluster.
    # These are perfect matches hence return 1.0.
    if (
        ref_classes.shape[0] == est_classes.shape[0] == 1
        or ref_classes.shape[0] == est_classes.shape[0] == 0
        or (ref_classes.shape[0] == est_classes.shape[0] == len(reference_indices))
    ):
        return 1.0

    contingency = _contingency_matrix(reference_indices, estimated_indices)

    # Compute the ARI using the contingency data
    sum_comb_c = sum(
        scipy.special.comb(n_c, 2, exact=1) for n_c in contingency.sum(axis=1)
    )
    sum_comb_k = sum(
        scipy.special.comb(n_k, 2, exact=1) for n_k in contingency.sum(axis=0)
    )

    sum_comb = sum(
        (scipy.special.comb(n_ij, 2, exact=1) for n_ij in contingency.flatten())
    )
    prod_comb = (sum_comb_c * sum_comb_k) / float(scipy.special.comb(n_samples, 2))
    mean_comb = (sum_comb_k + sum_comb_c) / 2.0
    return (sum_comb - prod_comb) / (mean_comb - prod_comb)
//  [python]
def _estimate_shapley_values_exact(
    set_func: Callable[[np.ndarray], Union[float, np.ndarray]], num_players: int, n_jobs: int
) -> np.ndarray:
    """Following Eq. (2) in
    Janzing, D., Minorics, L., & Bloebaum, P. (2020).
    Feature relevance quantification in explainable AI: A causal problem.
    In International Conference on Artificial Intelligence and Statistics (pp. 2907-2916). PMLR."""
    all_subsets = [tuple(subset) for subset in itertools.product([0, 1], repeat=num_players)]

    with Parallel(n_jobs=n_jobs) as parallel:
        subset_to_result_map = _evaluate_set_function(set_func, all_subsets, parallel)

    def compute_subset_weight(length: int) -> float:
        return 1 / (num_players * s.comb(num_players - 1, length))

    subset_weight_cache = {}

    shapley_values = [None] * num_players
    subsets_missing_one_player = np.array(list(itertools.product([0, 1], repeat=num_players - 1)))
    for player_index in range(num_players):
        subsets_with_player = [
            tuple(subset) for subset in np.insert(subsets_missing_one_player, player_index, 1, axis=1)
        ]
        subsets_without_player = [
            tuple(subset) for subset in np.insert(subsets_missing_one_player, player_index, 0, axis=1)
        ]

        for i in range(len(subsets_with_player)):
            subset_length = int(np.sum(subsets_without_player[i]))
            if subset_length not in subset_weight_cache:
                subset_weight_cache[subset_length] = compute_subset_weight(subset_length)

            weighted_diff = subset_weight_cache[subset_length] * (
                subset_to_result_map[subsets_with_player[i]] - subset_to_result_map[subsets_without_player[i]]
            )
            # For estimating Shapley values for multiple samples (e.g. in feature relevance) and the number of samples
            # is unknown beforehand.
            if shapley_values[player_index] is None:
                shapley_values[player_index] = weighted_diff
            else:
                shapley_values[player_index] += weighted_diff

    return np.array(shapley_values).T