aspcore.matrices.apply_blockwise

aspcore.matrices.apply_blockwise(mat, func, out_shape, *args, num_blocks=None, block_size=None, separate_axis=False, dtype=<class 'float'>, **kwargs)

Applies the same function to each block in the block matrix mat. Assumes that the matrix mat is square, and each block is square.

Either num_blocks or block_size must be supplied.

Parameters:

• mat (ndarray of shape (a, a))

• func (function) – function to be applied to each block, should return a matrix of shape out_shape

• out_shape (tuple of ints or 'same') – if ‘same’, the output will be the same size as the input blocks

• num_blocks (int, optional) – number of blocks in each dimension

• block_size (int, optional) – size of each block

• separate_axis (bool, optional) –

  if True, out_shape can be a scalar value or a tuple of any length

  the output will be of shape (*out_shape, num_blocks, num_blocks)

  if False, out_shape must be length-2 tuple

  the output will be of shape (num_blocks * out_shape[0], num_blocks * out_shape[1])

• dtype (type)

• kwargs (dict) – passed to func

Returns:

out – or ndarray of shape (*out_shape, num_blocks, num_blocks)

Return type:

ndarray of shape (num_blocks * out_shape[0], num_blocks * out_shape[1])