Sympy lambdify ImmutableDenseMatrix with numexpr

Question

I try to accelerate the evaluation of a MutableDenseMatrix using lambdify. It works with the module 'numpy'. 'Numexpr' should be faster (as I need the evaluation to solve a large optimization problem).

A smaller example of what I am trying to do is given by

from sympy import symbols, cos, Matrix, lambdify

a11, a12, a21, a22, b11, b12, b21, b22, u = symbols("a11 a12 a21 a22 b11 b12 b21 b22 u")
A = Matrix([[a11, a12], [a21, a22]])
B = Matrix([[b11, b12], [b21, b22]])
expr = A * (B ** 2) * cos(u) + A ** (-3 / 2)
f = lambdify((A, B), expr, modules='numexpr')

It raises the error

TypeError: numexpr cannot be used with ImmutableDenseMatrix

Is there a way to use lambdify for DenseMatrices? Or another idea how to speed up the evaluation?

Thanks in advance!

At first it may make sense to see the actual code which sympy exports. https://stackoverflow.com/a/55760092/4045774 `inspect.getsource(f)` After that there are many ways to optimize the implementation manually (Cython, Numba, Numexpr) or maybe also by a better vectorization strategy (Numpy) — max9111, Sep 07 '20 at 12:22
The code is way too long to show it here. It returns an ndarray of size (2,1). The two entries are a long mix containing +, -, *, /, ** and sqrt. I hope this helps. — Franzi, Sep 08 '20 at 14:08
A small example (at least with the same input's and outputs) or a full one (pastebin ?) would be good. It looks like quite simple to optimize using numba or cython, but there may be some manual editing and a few optimizations necessary. — max9111, Sep 08 '20 at 15:24
Here it is: https://pastebin.pl/view/31defe71 x,y and z are integer scalars. s0,su,sv,suu,suv,svv,u and v are ndarrays of size(383,552). Thank you! — Franzi, Sep 09 '20 at 10:15
This is a working numexpr version including benchmark (3700x speedup) https://pastebin.pl/view/bb190db6 . But it would be a good idea to add an example with a smaller MutableDenseMatrix to write an answer on that. It might be useful for others. I also found this sympy.org/scipy-2017-codegen-tutorial but in this example a text editing is enough. You also had code errors in your pastebin (The function header is in between at line 3118) — max9111, Sep 09 '20 at 20:19
Thank you very much. I edit the post, so its easier to give a precise answer. Would be really nice if you could explain what you did to get the working numexpr version. — Franzi, Sep 10 '20 at 11:53

score 1 · Accepted Answer · answered Sep 24 '20 at 17:29

A possible solution using numexpr is to evaluate every matrix expression on it's own. The following Code should output a python function which evaluates all matrix expressions using Numexpr.

Numexpr with Matrices

import numpy as np
import sympy

def lambdify_numexpr(args,expr,expr_name):
    from sympy.printing.lambdarepr import NumExprPrinter as Printer
    printer = Printer({'fully_qualified_modules': False, 'inline': True,'allow_unknown_functions': False})

    s=""
    s+="import numexpr as ne\n"
    s+="from numpy import *\n"
    s+="\n"

    #get arg_names
    arg_names=[]
    arg_names_str=""
    for i in range(len(args)):
        name=[ k for k,v in globals().items() if v is args[i]][0]
        arg_names_str+=name
        arg_names.append(name)

        if i< len(args)-1:
            arg_names_str+=","

    #Write header
    s+="def "+expr_name+"("+arg_names_str+"):\n"

    #unroll array
    for ii in range(len(args)):
        arg=args[ii]
        if arg.is_Matrix:
            for i in range(arg.shape[0]):
                for j in range(arg.shape[1]):
                    s+="    "+ str(arg[i,j])+" = " + arg_names[ii]+"["+str(i)+","+str(j)+"]\n"

    s+="    \n"
    #If the expr is a matrix
    if expr.is_Matrix:
        #write expressions
        for i in range(len(expr)):
            s+="    "+ "res_"+str(i)+" = ne."+printer.doprint(expr[i])+"\n"
            s+="    \n"

        res_counter=0
        #write array
        s+="    return concatenate(("
        for i in range(expr.shape[0]):
            s+="("
            for j in range(expr.shape[1]):
                s+="res_"+str(res_counter)+","
                res_counter+=1
            s+="),"
        s+="))\n"

    #If the expr is not a matrix
    else:
        s+="    "+ "return ne."+printer.doprint(expr)+"\n"
    return s

Sympy lambdify ImmutableDenseMatrix with numexpr

1 Answers1