geossm.utils.block_diag_3D

geossm.utils.block_diag_3D(*arrs)

Stack 3D arrays into a single 3D block-diagonal array.

Generalizes scipy.linalg.block_diag() to a stack of matrices that share a common third (e.g. time) axis: each input array of shape (n_i, m_i, p) becomes one diagonal block of an output of shape (sum(n_i), sum(m_i), p), zero-padded elsewhere, for every index along the shared p axis independently. This is used to combine several latent factors’ (or formulas’/blocks’) per-timestep system matrices (e.g. loading/basis matrices) into one block-diagonal matrix per timestep.

Parameters:

*arrs (numpy.ndarray) – Two or more 3D arrays of shape (n_i, m_i, p). The first two dimensions may differ between arrays; the third dimension p must be the same for all of them.

Returns:

numpy.ndarray – Block-diagonal array of shape (sum(n_i), sum(m_i), p) and dtype float64.