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

Code before:
                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):
    """
    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
    if K.is_EXRAW and op(zero, b) != zero:
        Csdm = sdm_scalar_mul_exraw(A, b, op)
    else:
        Csdm = unop_dict(A, lambda aij: op(aij, b))
    return Csdm


def sdm_scalar_mul_exraw(A, b, op):
    zero = A.domain.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:
                Ci[j] = Cij
    return Csdm


def sdm_irref(A):

Review comment (suggestion):
The `sdm_*` functions are supposed to just work with dicts. If needed the domain can be passed in as a separate argument but it should not be assumed that `A` has a domain attribute because in general a dict does not have that attribute.

Code after:
                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:
                Ci[j] = Cij
    return Csdm


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