PETScVector.set_local extremely slow

[Zane Jakobs]

Hi dolfin devs,

I have some code that, among other things, sets the local values of a tensor Function from a numpy array. I’m calling

my_dolfin_fun.vector().set_local(my_np_array.flatten()),

and the actual set_local call takes 1.2-1.3 seconds when my_np_array contains 300 2-by-2 matrices (i.e. has shape (300,2,2)). get_local takes $$2.2\cdot 10^{-5}$$ seconds, and the calls to my_dolfin_fun.vector() and my_np_array.flatten() both take order of $$10^{-5}$$ seconds or less, so I can be certain that set_local() is the problem. I’m very surprised that this is happening, given how simple the C++ code for PETScVector.set_local() is, and how fast the PETSc functions called inside of it usually are when I call them in C code. Any idea what’s going wrong here?

‌

EDIT: The issue was that the data was in a jax array, not in a numpy array. Explicitly converting fixed the problem.

[Zane Jakobs]

The issue was that the data was in a jax array, not in a numpy array. Explicitly converting fixed the problem.