def solve_linear(A, b, iterative_method=False):
    """ Master function to call the others """

    if iterative_method and iterative_method is not None:
        # if iterative solver string is supplied, use that method
        return _solve_iterative(A, b, iterative_method=iterative_method)
    elif iterative_method and iterative_method is None:
        # if iterative_method is supplied as None, use the default
        return _solve_iterative(A, b, iterative_method=DEFAULT_ITERATIVE_METHOD)
    else:
        # otherwise, use a direct solver
        return _solve_direct(A, b) def make_sparse(entries, indices, shape):
    """Construct a sparse csr matrix
    Args:
      entries: numpy array with shape (M,) giving values for non-zero
        matrix entries.
      indices: numpy array with shape (2, M) giving x and y indices for
        non-zero matrix entries.
      shape: shape of resulting matrix
    Returns:
      sparse, complex, matrix with specified values
    """  
    coo = sp.coo_matrix((entries, indices), shape=shape, dtype=npa.complex128)
    return coo.tocsr() def transpose_indices(indices):
    # returns the transposed indices for transpose sparse matrix creation
   return npa.flip(indices, axis=0) def make_rand(N):
    # makes a random vector of size N with elements between -0.5 and 0.5
    return npa.random.random(N) - 0.5 def make_rand_complex(N):
    # makes a random complex-valued vector of size N with re and im parts between -0.5 and 0.5
    return make_rand(N) + 1j * make_rand(N) def make_rand_indeces(N, M):
    # make M random indeces into an NxN matrix
    return npa.random.randint(low=0, high=N, size=(2, M)) def make_rand_sparse(N, M):
    # make a random sparse matrix of shape '(N, N)' and 'M' non-zero elements
    entries, indices = make_rand_entries_indices(N, M)
    return make_sparse(entries, indices, shape=(N, N)) def der_num(fn, arg, index, delta):
    # numerical derivative of `fn(arg)` with respect to `index` into arg and numerical step size `delta`
    arg_i_for  = arg.copy()
    arg_i_back = arg.copy()
    arg_i_for[index] += delta / 2
    arg_i_back[index] -= delta / 2
    df_darg = (fn(arg_i_for) - fn(arg_i_back)) / delta
    return df_darg def grad_num(fn, arg, delta=1e-6):
    # take a (complex) numerical gradient of function 'fn' with argument 'arg' with step size 'delta'
    N = arg.size
    grad = npa.zeros((N,), dtype=npa.complex128)
    f0 = fn(arg)
    for i in range(N):
        grad[i] = der_num(fn, arg, i, delta)        # real part
        grad[i] += der_num(fn, arg, i, 1j * delta)  # imaginary part
    return grad def get_entries_indices(csr_matrix):
    # takes sparse matrix and returns the entries and indeces in form compatible with 'make_sparse'
    shape = csr_matrix.shape
    coo_matrix = csr_matrix.tocoo()
    entries = csr_matrix.data
    cols = coo_matrix.col
    rows = coo_matrix.row
    indices = npa.vstack((rows, cols))
    return entries, indices def make_IO_matrices(indices, N):
    """ Makes matrices that relate the sparse matrix entries to their locations in the matrix
            The kth column of I is a 'one hot' vector specifing the k-th entries row index into A
            The kth column of J is a 'one hot' vector specifing the k-th entries columnn index into A
            O = J^T is for notational convenience.
            Armed with a vector of M entries 'a', we can construct the sparse matrix 'A' as:
                A = I @ diag(a) @ O
            where 'diag(a)' is a (MxM) matrix with vector 'a' along its diagonal.
            In index notation:
                A_ij = I_ik * a_k * O_kj
            In an opposite way, given sparse matrix 'A' we can strip out the entries `a` using the IO matrices as follows:
                a = diag(I^T @ A @ O^T)
            In index notation:
                a_k = I_ik * A_ij * O_kj
    """
    M = indices.shape[1]                                 # number of indices in the matrix
    entries_1 = npa.ones(M)                              # M entries of all 1's
    ik, jk = indices                                     # separate i and j components of the indices
    indices_I = npa.vstack((ik, npa.arange(M)))          # indices into the I matrix
    indices_J = npa.vstack((jk, npa.arange(M)))          # indices into the J matrix
    I = make_sparse(entries_1, indices_I, shape=(N, M))  # construct the I matrix
    J = make_sparse(entries_1, indices_J, shape=(N, M))  # construct the J matrix
    O = J.T                                              # make O = J^T matrix for consistency with my notes.
    return I, O
import autograd.numpy as npa
import scipy.sparse as sp
import autograd as ag
    import ceviche    # use the ceviche wrapper for autograd derivatives
from .solvers import solve_linear
from .utils import (make_sparse, transpose_indices, make_rand, make_rand_complex, make_rand_indeces,
                    make_rand_sparse, der_num, grad_num, get_entries_indices, make_IO_matrices)
    Returns:
      1d numpy array corresponding to the result (b) of A * x = b.
    """
    N = x.size
    A = make_sparse(entries, indices, shape=(N, N))
    return A.dot(x)

def grad_sp_mult_entries_reverse(ans, entries, indices, x):
    # x^T @ dA/de^T @ v => the outer product of x and v using the indices of A
    ia, ja = indices
    def vjp(v):
        return v[ia] * x[ja]
    return vjp

def grad_sp_mult_x_reverse(b, entries, indices, x):
    # dx/de^T @ A^T @ v => multiplying A^T by v
    indices_T = transpose_indices(indices)
    def vjp(v):
        return sp_mult(entries, indices_T, v)
    return vjp

ag.extend.defvjp(sp_mult, grad_sp_mult_entries_reverse, None, grad_sp_mult_x_reverse)

def grad_sp_mult_entries_forward(g, b, entries, indices, x):
    # dA/de @ x @ g => use `g` as the entries into A and multiply by x
    return sp_mult(g, indices, x)

def grad_sp_mult_x_forward(g, b, entries, indices, x):
    # A @ dx/de @ g -> simply multiply A @ g
    return sp_mult(entries, indices, g)

ag.extend.defjvp(sp_mult, grad_sp_mult_entries_forward, None, grad_sp_mult_x_forward)


""" ========================== Sparse Matrix-Vector Solve =========================="""

@ag.primitive
def sp_solve(entries, indices, b):
    """ Solve a sparse matrix (A) with source (b)
    Args:
      entries: numpy array with shape (num_non_zeros,) giving values for non-zero
        matrix entries.
      indices: numpy array with shape (2, num_non_zeros) giving x and y indices for
        non-zero matrix entries.
      b: 1d numpy array specifying the source.
    Returns:
      1d numpy array corresponding to the solution of A * x = b.
    Note: Calls a customizable solving function from ceviche.solvers, could add options to sp_solve() here eventually
    """
    N = b.size
    A = make_sparse(entries, indices, shape=(N, N))
    # calls a customizable solving function from ceviche.solvers, could add options to sp_solve() here eventually
    return solve_linear(A, b)

def grad_sp_solve_entries_reverse(x, entries, indices, b):
    # x^T @ dA/de^T @ A_inv^T @ -v => do the solve on the RHS, then take outer product with x using indices of A
    indices_T = transpose_indices(indices)
    i, j = indices
    def vjp(v):
        adj = sp_solve(entries, indices_T, -v)
        return adj[i] * x[j]
    return vjp

def grad_sp_solve_b_reverse(ans, entries, indices, b):
    # dx/de^T @ A_inv^T @ v => do the solve on the RHS and you're done.
    indices_T = transpose_indices(indices)
    def vjp(v):
        return sp_solve(entries, indices_T, v)
    return vjp

ag.extend.defvjp(sp_solve, grad_sp_solve_entries_reverse, None, grad_sp_solve_b_reverse)

def grad_sp_solve_entries_forward(g, x, entries, indices, b):
    # -A_inv @ dA/de @ A_inv @ b @ g => insert x = A_inv @ b and multiply with g using A indices.  Then solve as source for A_inv.
    forward = sp_mult(g, indices, x)
    return sp_solve(entries, indices, -forward)

def grad_sp_solve_b_forward(g, x, entries, indices, b):
    # A_inv @ db/de @ g => simply solve A_inv @ g
    return sp_solve(entries, indices, g)

ag.extend.defjvp(sp_solve, grad_sp_solve_entries_forward, None, grad_sp_solve_b_forward)


""" ==========================Sparse Matrix-Sparse Matrix Multiplication ========================== """

@ag.primitive
def spsp_mult(entries_a, indices_a, entries_x, indices_x, N):
    """ Multiply a sparse matrix (A) by a sparse matrix (X) A @ X = B
    Args:
      entries_a: numpy array with shape (num_non_zeros,) giving values for non-zero
        matrix entries into A.
      indices_a: numpy array with shape (2, num_non_zeros) giving x and y indices for
        non-zero matrix entries into A.
      entries_x: numpy array with shape (num_non_zeros,) giving values for non-zero
        matrix entries into X.
      indices_x: numpy array with shape (2, num_non_zeros) giving x and y indices for
        non-zero matrix entries into X.
      N: all matrices are assumed of shape (N, N) (need to specify because no dense vector supplied)
    Returns:
      entries_b: numpy array with shape (num_non_zeros,) giving values for non-zero
        matrix entries into the result B.
      indices_b: numpy array with shape (2, num_non_zeros) giving i, j indices for
        non-zero matrix entries into the result B.      
    """
    A = make_sparse(entries_a, indices_a, shape=(N, N))
    X = make_sparse(entries_x, indices_x, shape=(N, N))
    B = A.dot(X)
    entries_b, indices_b = get_entries_indices(B)
    return entries_b, indices_b

def grad_spsp_mult_entries_a_reverse(b_out, entries_a, indices_a, entries_x, indices_x, N):
    """ For AX=B, we want to relate the entries of A to the entries of B.
        The goal is to compute the gradient of the output entries with respect to the input.
        For this, though, we need to convert into matrix form, do our computation, and convert back to the entries.
        If you write out the matrix elements and do the calculation, you can derive the code below, but it's a hairy derivation.
    """

    # make the indices matrices for A
    _, indices_b = b_out

    Ia, Oa = make_IO_matrices(indices_a, N)