Repository: sympy/sympy
File: sympy/polys/matrices/sdm.py (Python)

Code before:
                        Ci = C.get(i, {})
                        Cij = Ci.get(j, zero) + Aik * Bkj
                        if Cij != zero:
                            Ci[j] = Cij
                            C[i] = Ci
                        else:
                            Ci.pop(j, None)
                            if Ci:
                                C[i] = Ci
                            else:
                                C.pop(i, None)

    return C


def sdm_scalar_mul(A, b, op):
    """
    Handles special cases like 0 * oo -> nan by creating a dense result
    when necessary. For all other cases, it uses the fast sparse approach.
    """

    K = A.domain
    zero = K.zero
    zero_prod = op(zero, b)

    if K.is_EXRAW and (zero_prod != zero):
        m, n = A.shape
        Csdm = {i: dict.fromkeys(range(n), zero_prod) for i in range(m)}
        for i, Ai in A.items():
            Ci = Csdm[i]
            for j, Aij in Ai.items():
                Cij = op(Aij, b)
                if Cij == zero:
                    del Ci[j]
                else:
                    Ci[j] = Cij
        return A.new(Csdm, A.shape, K)

    Csdm = unop_dict(A, lambda aij: op(aij, b))
    return A.new(Csdm, A.shape, K)


def sdm_irref(A):
    """RREF and pivots of a sparse matrix *A*.

    Compute the reduced row echelon form (RREF) of the matrix *A* and return a
    list of the pivot columns. This routine does not work in place and leaves
    the original matrix *A* unmodified.

    The domain of the matrix must be a field.


Review comment (suggestion):
This should check `K.is_EXRAW` first and then call a separate function rather than computing `zero_prod` up front. We want to keep the fast path for non EXRAW domains as fast as possible.

Code after:
            C[i] = Ci

    # Find all infinities in B
    for k, Bk in B.items():
        for j, Bkj in Bk.items():
            if zero * Bkj != zero:
                for i in range(m):
                    Aik = A.get(i, {}).get(k, zero)
                    # If Aik is not zero then this was handled above
                    if Aik == zero:
                        Ci = C.get(i, {})
                        Cij = Ci.get(j, zero) + Aik * Bkj
                        if Cij != zero:
                            Ci[j] = Cij
                            C[i] = Ci
                        else:
                            Ci.pop(j, None)
                            if Ci:
                                C[i] = Ci
                            else:
                                C.pop(i, None)

    return C


def sdm_scalar_mul(A, b, op, K):
    """
    Handles special cases like 0 * oo -> nan by creating a dense result
    when necessary. For all other cases, it uses the fast sparse approach.
    """

    zero = K.zero
    if K.is_EXRAW and op(zero, b) != zero:
        Csdm = sdm_scalar_mul_exraw(A, b, op, K)
    else:
        Csdm = unop_dict(A, lambda aij: op(aij, b))
    return Csdm


def sdm_scalar_mul_exraw(A, b, op, K):
    zero = K.zero
    zero_prod = op(zero, b)
    m, n = A.shape
    Csdm = {i: dict.fromkeys(range(n), zero_prod) for i in range(m)}
    for i, Ai in A.items():
        Ci = Csdm[i]
        for j, Aij in Ai.items():
            Cij = op(Aij, b)
            if Cij == zero:
                del Ci[j]
            else: