Search code examples
pythonfunctionscipyarguments

How to Call the Same Argument N Times in Another Function in Python


I have a function that can take a variable number of arguments. I want another function to call that function with one argument repeated n times. The function with variable arguments is scipy.linalg.block_diag, which creates a block diagonal matrix with its inputs along the diagonal. The function calling block_diag has a matrix and the number of times the matrix should be repeated along the diagonal as arguments. Here is what I have so far:

from scipy.linalg import block_diag

mat = [[0.22463041,  0.97444404,  0.], [0.97444404, -0.22463041,  0.], [ 0., 0., 1.]]
NDim = 7

def block_diagonal_matrix(matrix, NDimen):
    """
    This function generates a block diagonal matrix by repeating the input 
    matrix along the diagonal.

    Args:
        matrix: The input matrix.
        NDimen: The number of times the matrix should be repeated along the 
        diagonal.

    Returns:
        BlockDiagMat: The resulting block diagonal matrix obtained by 
        repeating the input matrix along the diagonal.
    """


    BlockDiagMat = block_diag(matrix, matrix, matrix, matrix, matrix, matrix, matrix)
    return BlockDiagMat

block_diagonal_matrix(mat, NDim)

By manually repeating matrix 7 (NDim) times when calling block_diag, I obtain the correct block diagonal matrix. However, I would like a way to have the block_diagonal_matrix function call matrix the appropriate number of times.

I tried calling block_diag like this within the function:

BlockDiagMat = block_diag(matrix*NDimen)

However, the resulting block diagonal matrix did not have the correct inner dimensions (expecting len(BlockDiagMat [0]) to be 21, got 3).

I also tried to make the matrix a list first:

BlockDiagMat = block_diag([matrix]*NDimen)

But scipy.linalg.block_diag did not support this: ValueError: arguments in the following positions have dimension greater than 2: [0]

I then tried to flatten with numpy.ndarray.flatten, but the resulting matrix did not have the correct dimensions again.

This all boils down to calling the matrix n times in scipy.linalg.block_diag using block_diagonal_matrix.

Thanks in advance for your help.


Solution

  • Try this:

    BlockDiagMat = block_diag(*[matrix]*NDimen)
    

    In general to repeat an argument n times into a function:

    func(*[arg]*n)
    

    First you create a list containing the argument [arg]

    Then you multiply it to have the argument the correct number of times [arg]*n

    Then you unpack this list into the function so that the arguments are actually given to the function as multiple arguments, rather than a single list *[arg]*n