ENH: set data keys as first positional arguments

[redeboer]

The positional arguments of the internal function of a ParametrizedBackendFunction are now ordered: args for the parameters come after args for the data sample over which the __call__() is evaluated. This is done so that one can compute the gradient of the backend function, e.g. with jax.jacfwd():

from tensorwaves.function import ParametrizedBackendFunction
from tensorwaves.interface import DataSample
data: DataSample
intensity_func: ParametrizedBackendFunction

import inspect
signature = inspect.signature(intensity_func.function)
arg_names = tuple(signature.parameters)
data_columns = {
    arg: data[key]
    for arg, key in zip(arg_names, intensity_func.argument_order)
    if key in data
}
parameter_values = {
    arg: complex(intensity_func.parameters[key]).real
    for arg, key in zip(arg_names, intensity_func.argument_order)
    if key in intensity_func.parameters
}

import jax
from jax.tree_util import Partial
func_with_data_inserted = Partial(intensity_func.function, *data_columns.values())
gradient = jax.jacfwd(func_with_data_inserted, argnums=range(len(parameter_values)))
gradient_values = gradient(*parameter_values.values())

We used Partial integrate the data columns into the function. It may be possible to integrate that into the gradient_creator() in order to address e.g. #442.