Repository: scipy/scipy
File: scipy/stats/tests/test_multivariate.py (Python)

Code before:


def _random_covariance(dim, evals, rng, singular=False):
    # Generates random covariance matrix with dimensionality `dim` and
    # eigenvalues `evals` using provided Generator `rng`. Randomly sets
    # some evals to zero if `singular` is True.
    A = rng.random((dim, dim))
    A = A @ A.T
    _, v = np.linalg.eigh(A)
    if singular:
        zero_eigs = rng.normal(size=dim) > 0
        evals[zero_eigs] = 0
    cov = v @ np.diag(evals) @ v.T
    return cov


def _sample_orthonormal_matrix(n):
    rng = np.random.default_rng(9086764251)
    M = rng.standard_normal((n, n))
    u, s, v = scipy.linalg.svd(M)
    return u

def marginal_pdf(X, X_ndim, dimensions, x):
    # First sort input data based on order of dimensions
    dimensions = np.asarray(dimensions)
    dimensions[dimensions < 0] = X_ndim + dimensions[dimensions < 0]
    dim_sort_idx = dimensions.argsort()
    x = x[:, dim_sort_idx]
    
    i_marginalize = np.ones(X_ndim, dtype=bool)
    i_marginalize[dimensions] = False

    def g(z):
        y = np.empty((z.shape[0], x.shape[0], X_ndim))
        y[..., i_marginalize] = z[:, np.newaxis, :]
        y[..., ~i_marginalize] = x
        return X.pdf(y)

    inf = np.full(X_ndim - len(dimensions), np.inf)
    return cubature(g, -inf, inf).estimate

@dataclass
class MVNProblem:
    """Instantiate a multivariate normal integration problem with special structure.

    When covariance matrix is a correlation matrix where the off-diagonal entries
    ``covar[i, j] == lambdas[i]*lambdas[j]`` for ``i != j``, then the multidimensional
    integral reduces to a simpler univariate integral that can be numerically integrated
    easily.

    The ``generate_*()`` classmethods provide a few options for creating variations

Review comment (suggestion):
nitpick:

```suggestion
    dimensions[dimensions < 0] += X_ndim
```

Code after:


def _random_covariance(dim, evals, rng, singular=False):
    # Generates random covariance matrix with dimensionality `dim` and
    # eigenvalues `evals` using provided Generator `rng`. Randomly sets
    # some evals to zero if `singular` is True.
    A = rng.random((dim, dim))
    A = A @ A.T
    _, v = np.linalg.eigh(A)
    if singular:
        zero_eigs = rng.normal(size=dim) > 0
        evals[zero_eigs] = 0
    cov = v @ np.diag(evals) @ v.T
    return cov


def _sample_orthonormal_matrix(n):
    rng = np.random.default_rng(9086764251)
    M = rng.standard_normal((n, n))
    u, s, v = scipy.linalg.svd(M)
    return u

def marginal_pdf(X, X_ndim, dimensions, x):
    """Integrate marginalized dimensions of multivariate
    probability distribution to calculate the marginalized
    distribution.
    """
    # Sort input data based on order of dimensions
    dimensions = np.asarray(dimensions)
    dimensions[dimensions < 0] += X_ndim
    dim_sort_idx = dimensions.argsort()
    x = x[:, dim_sort_idx]
    
    i_marginalize = np.ones(X_ndim, dtype=bool)
    i_marginalize[dimensions] = False

    def g(z):
        y = np.empty((z.shape[0], x.shape[0], X_ndim))
        y[..., i_marginalize] = z[:, np.newaxis, :]
        y[..., ~i_marginalize] = x
        return X.pdf(y)

    inf = np.full(X_ndim - len(dimensions), np.inf)
    return cubature(g, -inf, inf).estimate

@dataclass
class MVNProblem:
    """Instantiate a multivariate normal integration problem with special structure.

    When covariance matrix is a correlation matrix where the off-diagonal entries
    ``covar[i, j] == lambdas[i]*lambdas[j]`` for ``i != j``, then the multidimensional