diag

Contents

diag#

class brainunit.math.diag(v, k=0, unit=Unit('1'))#

Extract a diagonal or construct a diagonal array.

If v is a 1-D array, diag constructs a 2-D array with v on the k-th diagonal. If v is a 2-D array, diag extracts the k-th diagonal and returns a 1-D array.

Parameters:

  • v (saiunit.Quantity | Array | ndarray | bool | number | bool | int | float | complex) – Input array. 1-D inputs produce a 2-D diagonal matrix; 2-D inputs have their k-th diagonal extracted.

  • k (int) – Diagonal in question. The default is 0. Use k > 0 for diagonals above the main diagonal and k < 0 for diagonals below.

  • unit (saiunit.Unit) – Unit of the returned Quantity. Ignored when v already carries a unit.

Returns:

out – The extracted diagonal or constructed diagonal array.

Return type:

saiunit.Quantity | Array

Examples

>>> import saiunit as u
>>> import jax.numpy as jnp
>>> u.math.diag(jnp.array([1.0, 2.0, 3.0]))
Array([[1., 0., 0.],
       [0., 2., 0.],
       [0., 0., 3.]], dtype=float32)
>>> u.math.diag(jnp.array([1.0, 2.0]), unit=u.meter)
Quantity([[1. 0.]
          [0. 2.]], "m")
Contents