Python 3: Sympy: Include list information to optimize lambdify

Question

I use lambdify to compile an expression which is a function of certain parameters. Each parameter has N points. So I need to evaluate the expression N times. The following shows a simplified example on how this is done.

import numpy as np
from sympy.parsing.sympy_parser import parse_expr
from sympy.utilities.lambdify import lambdify, implemented_function
from sympy import S, Symbol
from sympy.utilities.autowrap import ufuncify


def CreateMagneticFieldsList(dataToSave,equationString,DSList):

    expression  = S(equationString)
    numOfElements = len(dataToSave["MagneticFields"])

    #initialize the magnetic field output array
    magFieldsArray    = np.empty(numOfElements)
    magFieldsArray[:] = np.NaN

    lam_f = lambdify(tuple(DSList),expression,modules='numpy')
    try:
        for i in range(numOfElements):
            replacementList = np.zeros(len(DSList))


            for j in range(len(DSList)):
                replacementList[j] = dataToSave[DSList[j]][i]

            try:
                val = np.double(lam_f(*replacementList))

            except:
                val = np.nan
            magFieldsArray[i] = val
    except:
        print("Error while evaluating the magnetic field expression")
    return magFieldsArray


list={"MagneticFields":list(range(10000)), "Chx":list(range(10000))}

out=CreateMagneticFieldsList(list,"MagneticFields*5+Chx",["MagneticFields","Chx"])

print(out)

Is there a way to optimize this call further? Specifically, I mean is there a way to make lambdify include that I'm calculating for a list of points, so that the loop evalulation can be optimized?

Answer 1

Thanks to @asmeurer, he gave the idea on how to do it.

Since lambdify is compiled using numpy, then one could simply pass the lists as arguments! The following is a working example

#!/usr/bin/python3

import numpy as np
from sympy.parsing.sympy_parser import parse_expr
from sympy.utilities.lambdify import lambdify, implemented_function
from sympy import S, Symbol
from sympy.utilities.autowrap import ufuncify


def CreateMagneticFieldsListOpt(dataToSave,equationString,DSList):

    expression  = S(equationString)
    numOfElements = len(dataToSave["MagneticFields"])

    #initialize the magnetic field output array
    magFieldsArray    = np.empty(numOfElements)
    magFieldsArray[:] = np.NaN

    lam_f = lambdify(tuple(DSList),expression,modules='numpy')
    replacementList = [None]*len(DSList)

    for j in range(len(DSList)):
        replacementList[j] = np.array(dataToSave[DSList[j]])
    print(replacementList)

    magFieldsArray = np.double(lam_f(*replacementList))


    return magFieldsArray

list={"MagneticFields":[1,2,3,4,5],"ChX":[2,4,6,8,10]}

out=CreateMagneticFieldsListOpt(list,"MagneticFields*5+ChX",["MagneticFields","ChX"])

print(out)