adam_core.coordinates.jacobian module

adam_core.coordinates.jacobian.calc_jacobian(coords: ndarray, _func: Callable, in_axes: Hashable | None = (0,), out_axes: int | None = 0, **kwargs) → Array[source]

Calculate the jacobian for the given callable in D dimensions for every N coordinate.

Parameters:

  • coords (~numpy.ndarray (N, D)) – Coordinates that correspond to the input covariance matrices.

  • _func (function) – A function that takes a single coord (D) as input and return the transformed coordinate (D). See for example: adam_core.coordinates.transform._cartesian_to_spherical or adam_core.coordinates.transform._cartesian_to_keplerian.

  • in_axes (Optional[Hashable]) –

    An integer or None indicates which array axis to map over for all arguments (with None indicating not to map any axis), and a tuple indicates which axis to map for each corresponding positional argument. Axis integers must be in the range [-ndim, ndim) for each array, where ndim is the number of dimensions (axes) of the corresponding input array.

    From: https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html#jax.vmap

  • out_axes (Optional[int]) –

    An integer, None, or (nested) standard Python container (tuple/list/dict) thereof indicating where the mapped axis should appear in the output. All outputs with a mapped axis must have a non-None out_axes specification.

    From: https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html#jax.vmap

Returns:

jacobian – Array containing function partial derivatives for each coordinate.

Return type:

~numpy.ndarray (N, D, D)