brainevent.binary_jitumm

Contents

brainevent.binary_jitumm#

brainevent.binary_jitumm = <NameScope(brainevent.binary_jitumm)>#

Event-driven matrix-matrix product with a JIT uniform connectivity matrix.

Computes the product of a just-in-time generated sparse matrix with uniformly distributed weights and a binary event matrix B. Each column of B is treated as an independent event vector, and only non-zero entries contribute to the output.

The sparse matrix A of shape (m, n) is never materialized. Each entry A[i, j] is drawn from Uniform(w_low, w_high) with probability prob, seeded by seed.

Parameters:
  • w_low (Array | ndarray | Quantity | Number) – Lower bound of the uniform weight distribution. Scalar value, optionally with physical units (brainunit.Quantity).

  • w_high (Array | ndarray | Quantity | Number) – Upper bound of the uniform weight distribution. Must have the same dimension (units) as w_low.

  • prob (float) – Connection probability in [0, 1]. Determines the fraction of non-zero entries in the connectivity matrix.

  • B (Array | ndarray | Quantity | Number) – Input event matrix of shape (n, k) (if transpose=False) or (m, k) (if transpose=True). Can be boolean or floating-point; non-zero entries are treated as active events.

  • seed (int | None) – Random seed for reproducible connectivity patterns. If None, a random seed is generated at compile time.

  • shape (Tuple[int, int]) – Shape (m, n) of the logical connectivity matrix.

  • transpose (bool) – If True, compute A.T @ B instead of A @ B. Default is False.

  • corder (bool) – Memory layout order for the connectivity generation. True for C-order (row-major), False for Fortran-order (column-major). Default is True.

  • backend (str | None) – Computation backend. One of 'numba' or 'pallas'. If None, the default backend is used.

Returns:

Result matrix of shape (m, k) (if transpose=False) or (n, k) (if transpose=True). Carries the product of units from the weight and B if either has physical units.

Return type:

Array | ndarray | Quantity | Number

See also

binary_jitumv

Event-driven matrix-vector variant.

jitumm

Float (non-event) matrix-matrix variant.

Notes

The connectivity matrix A of shape (m, n) follows the model:

A[i, j] = U[i, j] * B_conn[i, j]

where U[i, j] ~ Uniform(w_low, w_high) and B_conn[i, j] ~ Bernoulli(prob) are independent, both determined by seed.

The event-driven matrix-matrix product computes:

result[i, j] = sum_{k : B[k, j] is active} A[i, k]

where “active” means True for boolean arrays or > 0 for float arrays. Each column of B is treated as an independent event vector. The full expansion is:

result[i, j] = sum_{k} U[i, k] * B_conn[i, k] * 1_{B[k, j] active}

When transpose=True, the operation becomes result = A^T @ B.

Examples

>>> import jax.numpy as jnp
>>> from brainevent._jit_uniform.binary import binary_jitumm
>>> B = jnp.array([[True, False], [False, True], [True, True],
...                [False, False], [True, False]])
>>> result = binary_jitumm(
...     0.1, 0.5, 0.2, B, seed=42,
...     shape=(3, 5), transpose=False, corder=True,
... )
>>> result.shape
(3, 2)