GaussNewton and LevenbergMarquardt on multidimensional arrays

[arunoruto]

I am currently trying to implement an inversion of Hapkes Anisotropic Multiple Scattering Approximation model. I implemented the AMSA code without problems and I am checking it against data obtained from a Matlab version of the code, so far so good! I even implemented a derivation function, which calculates the analytical derivation of the AMSA model.

Now I want to use those two functions to invert the model and obtain the albedo from a measured reflectance using a nonlinear least squares algorithm. This algorithm would have to go through each pixel and execute the least squares function to obtain the albedo and store it in an array. I did exactly that using a for-loop and so far it works, but it is reaaaaally slow. Here is the current implementation of the "inverted" model:

import jax
import jax.numpy as np
from jax.typing import ArrayLike
from jaxopt import GaussNewton, LevenbergMarquardt
from scipy.optimize import least_squares

from refmod.hapke import amsa
from refmod.hapke.amsa import amsa_derivative


def inverse_model(
    refl: ArrayLike,
    incidence_direction: ArrayLike,
    emission_direction: ArrayLike,
    surface_orientation: ArrayLike,
    phase_function: dict,
    b_n: ArrayLike,
    a_n: ArrayLike,
    roughness: float = 0,
    hs: float = 0,
    bs0: float = 0,
    hc: float = 0,
    bc0: float = 0,
) -> ArrayLike:
    albedo_recon = np.zeros_like(refl)
    s = np.prod(np.array(refl.shape))
    for k in range(s):
        print(f"{k / s* 100: .4f}%")  # if k % 20 == 0 else ""
        i = k % refl.shape[0]
        j = k // refl.shape[0]
        x0 = np.array([1 / 3])

        args = {
            "incidence_direction": incidence_direction[i, j, :],
            "emission_direction": emission_direction[i, j, :],
            "surface_orientation": surface_orientation[i, j, :],
            "phase_function": phase_function,
            "b_n": b_n,
            "a_n": a_n,
            "roughness": roughness,
            "hs": hs,
            "bs0": bs0,
            "hc": hc,
            "bc0": bc0,
            "refl_optimization": refl[i, j],
        }

        # lm = GaussNewton(
        lm = LevenbergMarquardt(
            residual_fun=amsa,
            # jac_fun=amsa_derivative,
            materialize_jac=True,
            # implicit_diff=False,
            # implicit_diff_solve=amsa_derivative,
            # jit=True,
        )
        lm_sol = lm.run(x0, **args).params
        albedo_recon = albedo_recon.at[i, j].set(lm_sol[0])

        # res = least_squares(
        #     amsa,
        #     x0,
        #     jac=amsa_derivative,
        #     method="lm",
        #     kwargs=args,
        # )
        # albedo_recon = albedo_recon.at[i, j].set(res.x[0])
        # x0 = res.x

    return albedo_recon

I am testing both scipy's least_squares and GN and LM from jaxopt. In this setup, scipy's least_squares wins.

Here are the runtimes across a 11x11 patch:

scipy:
Inverse AMSA: Mean +- std dev: 1.17 sec +- 0.03 sec
LM:
Inverse AMSA: Mean +- std dev: 42.7 sec +- 0.5 sec

GN takes even longer, but funny enough, is much closer to the actual value in testing!

Clearly this is not going to cut it and I have to somehow get rid of that for-loop. Both amsa and amsa_derivative are jitted and can take vectors as arguments, and this works really great for calculating reflectances from albedo maps ( about 30ms for a 1500x1500 image). How would I approach it here? Can I somehow supply the image and tell it, to optimize along a certain axis?

Small disclaimer: if you ask yourself now "why does he use a nonlinear least squares approach and not a root finder algorithms?" you are certainly correct! For this example a roots function would also suffice. But in the future, the function will take multiple images into account and try to find an albedo that fits all of them (albedo is in this case a constant and does not depend on the viewing geometry).