def diag(expr):
"""Extracts the diagonal from a matrix or makes a vector a diagonal matrix.
Parameters
----------
expr : Expression or numeric constant
A vector or square matrix.
Returns
-------
Expression
An Expression representing the diagonal vector/matrix.
"""
expr = AffAtom.cast_to_const(expr)
if expr.is_vector():
if expr.size[1] == 1:
return diag_vec(expr)
# Convert a row vector to a column vector.
else:
expr = reshape(expr, expr.size[1], 1)
return diag_vec(expr)
elif expr.size[0] == expr.size[1]:
return diag_mat(expr)
else:
raise ValueError("Argument to diag must be a vector or square matrix.")
def hstack(arg_list):
"""Horizontal concatenation of an arbitrary number of Expressions.
Parameters
----------
arg_list : list of Expression
The Expressions to concatenate.
"""
arg_list = [AffAtom.cast_to_const(arg) for arg in arg_list]
for idx, arg in enumerate(arg_list):
if arg.ndim == 0:
arg_list[idx] = arg.flatten()
return Hstack(*arg_list)
def diag(expr):
"""Extracts the diagonal from a matrix or makes a vector a diagonal matrix.
Parameters
----------
expr : Expression or numeric constant
A vector or square matrix.
Returns
-------
Expression
An Expression representing the diagonal vector/matrix.
"""
expr = AffAtom.cast_to_const(expr)
if expr.is_vector():
return diag_vec(vec(expr))
elif expr.ndim == 2 and expr.shape[0] == expr.shape[1]:
return diag_mat(expr)
else:
raise ValueError("Argument to diag must be a vector or square matrix.")
