How to compute derivatives with Lambdify and IndexedBase in sympy

I am investigating the suitability of SymPy for some of my projects and have come across an issue with the interaction between lambdify and IndexedBase.

In short, my applications make heavy use of functions that employ doubled summed array structures. I need to be able to compute both the function and the first through 3rd derivative of the function with respect to an array element.

My questions are thus:

How do I make lambdify work with diff?
How do I make a lambdify function where I can specify the index with respect to which I want the derivative?
How do I extend the above to a second order derivative with a different index (i.e. second order with respect to index i and j)?

Simplified example:

from sympy import IndexedBase, Idx, lambdify, Sum, diff
from numpy import array


i = Idx("i", range=(0,1))
j = Idx("j", range=(0,1))
n = IndexedBase("n")
coefficients = IndexedBase("C")

double_sum = Sum(Sum(n[i]*n[j]*coefficients[i,j],(i,i.lower,i.upper)),(j,j.lower,j.upper))
first_derivative = diff(double_sum, n[i])
second_derivative = diff(first_derivative, n[j])

test_function_1 = lambdify((n,coefficients),double_sum)
test_function_2 = lambdify((n,coefficients,i),first_derivative)
test_function_3 = lambdify((n,coefficients,i,j),second_derivative)

test_vector = array([1, 2])
test_coefficients = array([[1,1],[2,3]])

test_value_1 = test_function_1(test_vector,test_coefficients)
print(test_value_1)

test_value_2 = test_function_2(test_vector,test_coefficients,1)
print(test_value_2)

test_value_3 = test_function_3(test_vector,test_coefficients)
print(test_value_3)

Execution of this code yields the error:

File "<lambdifygenerated-2>", line 9, in _lambdifygenerated
File "<lambdifygenerated-2>", line 9, in <genexpr>
NameError: name 'KroneckerDelta' is not defined

Solution

Indexed expressions are useful, but derivatives of them are sometimes more complicated than they should be, and they tend to have issues with lambdify. Here is roughly equivalent code that does not use Indexed. The difference is that the size of array is declared upfront, which makes it possible to create explicit arrays of normal (non-indexed) symbols with symarray, manipulate those, and lambdify the expressions. I lambdified those in a way so that the first derivatives are returned as a 1-column matrix, and second derivatives as a square matrix (an alternative return type is below).

from sympy import symarray, Add, lambdify, Matrix
from numpy import array

i_upper = 2
j_upper = 2
n = symarray("n", i_upper)
coefficients = symarray("C", (i_upper, j_upper))

double_sum = Add(*[n[i]*n[j]*coefficients[i, j] for i in range(i_upper) for j in range(j_upper)])
first_derivative = Matrix(i_upper, 1, lambda i, j: double_sum.diff(n[i]))
second_derivative = Matrix(i_upper, j_upper, lambda i, j: double_sum.diff(n[i], n[j]))

params = list(n) + list(coefficients.ravel())
test_function_1 = lambdify(params, double_sum)
test_function_2 = lambdify(params, first_derivative)
test_function_3 = lambdify(params, second_derivative)

test_vector = array([1, 2])
test_coefficients = array([[1, 1], [2, 3]])
test_params = list(test_vector) + list(test_coefficients.ravel())    
test_value_1 = test_function_1(*test_params)
test_value_2 = test_function_2(*test_params)
test_value_3 = test_function_3(*test_params)

The test values are 19, [[8], [15]] and [[2, 3], [3, 6]] respectively.

Alternatively, the functions can return nested lists:

first_derivative = [double_sum.diff(n[i]) for i in range(i_upper)]
second_derivative = [[double_sum.diff(n[i], n[j]) for j in range(i_upper)] for i in range(i_upper)]

or NumPy arrays:

first_derivative = array([double_sum.diff(n[i]) for i in range(i_upper)])
second_derivative = array([[double_sum.diff(n[i], n[j]) for j in range(i_upper)] for i in range(i_upper)])

Limitations of this approach: (a) must know the number of symbols at the time of forming the expressions; (b) cannot accept indexes i or j as parameters of the lambdified function.