Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ENH: set data keys as first positional arguments #488

Merged
merged 5 commits into from
Nov 8, 2023
Merged

Conversation

redeboer
Copy link
Member

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.

@redeboer redeboer added the ⚙️ Enhancement Improvements and optimizations of existing features label May 26, 2023
@redeboer redeboer added this to the 0.4.11 milestone May 26, 2023
@redeboer redeboer self-assigned this May 26, 2023
@redeboer redeboer changed the title ENH: move data arguments to front for partial ENH: set data keys as first positional arguments May 26, 2023
This is required for defining a gradient over the parameters
The official webpages of Minuit2 appear to be unstable
@redeboer redeboer marked this pull request as ready for review November 8, 2023 13:32
@redeboer redeboer enabled auto-merge (squash) November 8, 2023 13:32
@redeboer redeboer merged commit aa36e66 into main Nov 8, 2023
20 of 21 checks passed
@redeboer redeboer deleted the order-function-args branch November 8, 2023 13:39
@redeboer redeboer mentioned this pull request Mar 6, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
⚙️ Enhancement Improvements and optimizations of existing features
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant