I have a non-linear forward model which computes a gray scale image for a per-pixel parameter w. I am also able to invert the model using scipys optimize functions. The only problem I am having currently is the size of the image is making this solution really slow… Like 7% of the pixels have been calculated in 40 minutes slow. I am looping over all pixels with a for-loop and applying the model pixelwise. I have tried the least_squares, leastsq, and root functions, but all of them are quite slow. I also implemented the derivative values, since I heard it could be faster by providing the jacobian matrix, but to no avail.

I thought the leastsq could be the solution, since it has the following statement in the description:

x = arg min(sum(func(y)**2,axis=0))
         y

If I reshaped my image as 1xn, it should minimize along the axis=0. But somehow I end up with an array of size 2 as my input argument y, and I wasn’t able to debug it.

I am especially looking at least square optimizations, since in the future, I would like to provide multiple images, which should be generated by the same parameter w (with other arguments varying).

Are there any libraries besides scipy, which would provide a more elegant approach to this problem?

I was also planning on trying a numba implementation of the root function, but it has its own problems. Or to use Googles JAX implementation of LM, but I have no real experience with JAX so far…