def index(i, j, n):
    """Function to return the index of an element of a 2D density matrix if flattened to a 1D array
    
    Parameters:
        i (State): To get the row number of the element to be selected
        j (State): To get the column number of the element to be selected
        n (int): dimension of the square array i.e. number of states in the system
    
    Returns:
        int: Index of the element in a 1D array
    
    Example:
        To get the index of rho_44 in a 4-level system call index(four, four, 4) where "four" is the variable which stores the State object labelled 4.
    """ 
    
    return ((i.label-1)*n+j.label)-1