JITCScalarMatrix

class brainevent.JITCScalarMatrix(weight, prob=None, seed=None, *, shape, corder=False, backend=None, buffers=None)

Base class for Just-In-Time Connectivity Scalar-weight matrices.

This abstract class serves as the foundation for sparse matrix representations that use homogeneous (scalar) weights with stochastic connectivity patterns. It stores a single weight value applied to all non-zero elements, along with connectivity probability and a random seed that determines the sparse structure.

Parameters:

• weight (WeightScalar or Tuple[WeightScalar, Prob, Seed]) – Either the homogeneous weight value for all non-zero elements, or a tuple containing (weight, prob, seed).

• prob (Prob, optional) – Connection probability determining matrix sparsity.

• seed (Seed, optional) – Random seed for reproducible sparse structure generation.

• shape (Tuple[int, int]) – The shape of the matrix as a tuple (rows, columns).

• corder (bool) – Memory layout order flag, by default False.

• backend (str | None) – The computation backend to use.

Returns:

A new scalar-weight JIT connectivity matrix instance.

Return type:

JITCScalarMatrix

Raises:

ValueError – If prob is not a scalar, is not finite, or is outside [0, 1].

See also

JITCScalarR

Row-oriented concrete subclass.

JITCScalarC

Column-oriented concrete subclass.

JITCMatrix

Parent class for all JIT connectivity matrices.

Notes

The matrix W is defined by a scalar weight w, a connection probability p, and a deterministic pseudo-random seed s. Each element is given by:

W[i, j] = w * B[i, j]

where B[i, j] ~ Bernoulli(p) is a binary mask whose realization is fully determined by the seed s. Specifically, the mask is generated by a deterministic hash-based PRNG that, for a given (i, j, s) triple, always produces the same binary outcome. This means:

• The same (weight, prob, seed, shape) always produces the identical matrix.

• The expected number of non-zeros per row is p * n_cols.

• The matrix is never materialized in memory; it is regenerated on-the-fly during each operation (matvec, matmat, todense).

The connection length parameter clen = 2 / p controls the average stride between successive non-zero entries in the sampling loop.

Examples

>>> from brainevent import JITCScalarR
>>> mat = JITCScalarR((0.5, 0.1, 42), shape=(100, 50))
>>> mat.weight   # 0.5
>>> mat.prob     # 0.1
>>> mat.seed     # 42

weight

The homogeneous weight value applied to all non-zero elements in the matrix. Can be a plain JAX array or a quantity with units.

Type:

Union[jax.Array, u.Quantity]

prob

Connection probability determining the sparsity of the matrix. Values range from 0 (no connections) to 1 (fully connected).

Type:

Union[float, jax.Array]

seed

Random seed controlling the specific pattern of connections. Using the same seed produces identical connectivity patterns.

Type:

Union[int, jax.Array]

shape

Tuple specifying the dimensions of the matrix as (rows, columns).

Type:

MatrixShape

corder

Flag indicating the memory layout order of the matrix. False (default) for Fortran-order (column-major), True for C-order (row-major).

Type:

bool

backend

The computation backend to use (e.g., 'numba', 'pallas'). If None, the default backend is selected automatically.

Type:

str or None

property data: Number | ndarray | Array | Quantity

Return the trainable weight of the matrix.

Only the trainable value parameter is exposed here. The structural parameters prob and seed are non-trainable and are therefore excluded. This property mirrors with_data(), which accepts exactly the value returned here, so mat.with_data(mat.data) round-trips.

Returns:

The homogeneous weight value for the non-zero elements.

Return type:

WeightScalar

See also

with_data

Rebuild the matrix from the value returned here.

dt2t(y_dim_arr, w_dim_arr)

Generate per-synapse weight * y[row] using the matrix parameters.

w_dim_arr is required by the DataRepresentation protocol and is not used. JITC scalar connectivity and weights are generated from this matrix’s own metadata, including weight, prob, seed, shape, corder, and backend.

Return type:

Array | Quantity

dt2t_transposed(y_dim_arr, w_dim_arr)

Generate per-synapse weight * y[col] using the matrix parameters.

w_dim_arr is required by the DataRepresentation protocol and is not used. JITC scalar connectivity and weights are generated from this matrix’s own metadata, including weight, prob, seed, shape, corder, and backend.

Return type:

Array | Quantity

property dtype

Get the data type of the matrix elements.

Returns:

The data type of the weight values in the matrix.

Return type:

dtype

Notes

This property inherits the dtype directly from the weight attribute, ensuring consistent data typing throughout operations involving this matrix.

tocsr()

Convert the sparse scalar-weight matrix to Compressed Sparse Row (CSR) format.

Generates the non-zero structure (data, indices, indptr) directly from the connectivity parameters using dedicated CPU/CUDA operators, without ever materializing the dense matrix. The resulting CSR reproduces exactly the same matrix as todense() for the active compute backend; every stored value equals the constant weight w.

Returns:

A CSR matrix with the same shape and values as todense(). The data type matches the weight, and physical units (brainunit.Quantity) are preserved on the stored values.

Return type:

CSR

See also

todense

Materialize the matrix as a dense array.

Notes

Generation uses a count pass followed by a fill pass; the number of stored elements is read back between the two passes, so tocsr is an eager-only conversion and cannot be traced under jax.jit. Peak memory is O(nnz) rather than O(rows * cols).

Examples

>>> from brainevent import JITCScalarR
>>> mat = JITCScalarR((1.5, 0.2, 42), shape=(10, 10))
>>> csr = mat.tocsr()
>>> csr.shape
(10, 10)

tree_flatten()

Flatten the matrix into leaves and auxiliary data for JAX pytree registration.

Returns:

A 2-tuple where the first element is a tuple of JAX-traceable leaves (weight, prob, seed) and the second element is a dict of static auxiliary data {'shape': ..., 'corder': ..., 'backend': ...}.

Return type:

tuple

See also

tree_unflatten

Reconstruct the matrix from flattened representation.

classmethod tree_unflatten(aux_data, children)

Reconstruct a matrix instance from flattened pytree data.

Parameters:

• aux_data (dict) – Dictionary of static auxiliary data containing 'shape', 'corder', and 'backend' keys.

• children (tuple) – Tuple of JAX-traceable leaves (weight, prob, seed).

Returns:

A reconstructed matrix instance with attributes restored from both children and aux_data.

Return type:

JITCScalarMatrix

See also

tree_flatten

Flatten the matrix for JAX pytree operations.

with_data(data)

Create a new matrix instance with updated weight, preserving all other structure.

Accepts exactly the value returned by data (the trainable weight), while keeping the same prob, seed, shape, corder, backend, and buffers. It is useful for updating the weight without changing the connectivity pattern.

Parameters:

data (Number | ndarray | Array | Quantity) – The new weight value to use. Must have the same shape and unit as the current weight.

Returns:

A new matrix instance of the same concrete type with the updated weight.

Return type:

JITCScalarMatrix

Raises:

AssertionError – If the provided weight has a different shape or unit than the current weight.

See also

data

Property returning the value accepted here.