Templates and Jacobians¶

This module provides functions to evaluate the Jacobian of a template function.

Templates¶

Template functions model the observed data. In this framework, we assume they are given in the Fourier domain as complex-valued functions of the model parameters and frequencies,

\[h_i(\theta)(f),\]

where \(i\) is the channel (observable) index, \(\theta\) is a set of model parameters, and \(f\) is the frequency vector.

The implementation of the template function is expected to be in JAX. It should be a function of the model parametrers as positional arguments and return an array of complex values for all frequencies and channels.

wejax.jacobian.Template¶

Template function.

Example

Below is an example of a template function that models a single channel with two parameters.

import jax.numpy as jnp
from jax import Array

freqs = jnp.linspace(0, 1, 10)

def template(a: float, b: float) -> Array:
    channel = a**2 * freqs + 42j * jnp.cos(a) * b**2
    return jnp.stack([channel], axis=-1)

template(1.0, 2.0)

This template can be vectorized for efficient computation of multiple sets of parameters using jax.vmap() (see JAX documentation).

from jax import vmap

a = jnp.array([1.0, 2.0, 3.0])
b = jnp.array([4.0, 5.0, 6.0])
vtemplate = vmap(template, in_axes=(0, 0), out_axes=0)
vtemplate(a, b)

Parameters:

*args (Array-like) – Model parameters as positional arguments.

Returns:

Template values.

The returned shape is often (..., Nf, Nc), with Nf is the number of frequencies, Nc is the number of channels. The ellipsis represents any number of additional dimensions.

Return type:

Complex-valued array

Jacobians¶

Given a template function \(h_i(\theta, f)\) that depends on a set of parameters \(\theta\) and frequencies \(f\), the Jacobian is a matrix of partial derivatives of the template with respect to the model parameters,

\[J_{ij}(\theta)(f) = \frac{\partial h_i}{\partial \theta_j}(\theta)(f).\]

The implementation of the Jacobian function is expected to be in JAX. It should be a function of the model parameters as positional arguments and return an array of complex values for all frequencies, channels, and parameters (along which the partial derivatives are computed).

wejax.jacobian.Jacobian¶

Jacobian function.

Example

import jax.numpy as jnp
from jax import Array

freqs = jnp.linspace(0, 1, 10)

def jac(a: float, b: float) -> Array:
    # Compute Jacobian elements
    jac_0 = jnp.stack([2 * a * freqs - 42j * b**2 * jnp.sin(a)], axis=-1)
    jac_1 = jnp.stack([84j * jnp.cos(a) * b * jnp.ones_like(freqs)], axis=-1)

    return jnp.stack([jac_0, jac_1], axis=-1)

jac(1.0, 2.0)

This Jacobian can be vectorized for efficient computation of multiple sets of parameters using jax.vmap() (see JAX documentation).

from jax import vmap

a = jnp.array([1.0, 2.0, 3.0])
b = jnp.array([4.0, 5.0, 6.0])
vjac = vmap(jac, in_axes=(0, 0), out_axes=0)
vjac(a, b)

Parameters:

*args (Array-like) – Model parameters as positional arguments.

Returns:

Template Jacobian values.

The returned shape is often (..., Nf, Nc, Np), with Nf is the number of frequencies, Nc is the number of channels, and Np the number of parameters. The ellipsis represents any number of additional dimensions.

Return type:

Complex-valued array

This module provides convenience functions to compute Jacobians using automatic differentiation and finite differences. It also provides a function to mix multiple Jacobians along the parameter axis.

Automatic differentiation¶

The autograd() function computes the Jacobian using automatic differentiation. This function uses JAX forward-mode automatic differentiation to compute the Jacobian of the template function.

Note

Vectorization

To vectorize the computation of the Jacobian for multiple sets of parameters, use jax.vmap() on the non-vectorized Jacobian function, to avoid computing partial derivatives for each parameter array element.

wejax.jacobian.autograd(template, argnums)[source]¶

Compute the Jacobian of a function using automatic differentiation.

This function use JAX forward-mode (column-wise jax.jacfwd()) automatic differentiation to compute the Jacobian of the template function.

Example

>>> from wejax.jacobian import autograd
>>> jac = autograd(template, argnums=(0, 1))
>>> jac(1.0, 2.0)
Array([[[[...]]]], dtype=complex64)

Parameters:

template (Callable[..., Array]) – Template function, implemented in JAX.
argnums (Sequence[int]) – Indices of the positional arguments to differentiate with respect to.

Returns:

Jacobian function.

Return type:

Callable[..., Array]

Finite differences¶

The findiff() function computes the Jacobian using finite differences. This function uses numerical differentiation to compute the Jacobian of the template function.

The function requires a set of steps for numerical differentiation. These can be either a scalar or an array with the same shape as the parameters to differentiate. If a scalar, the same step is used for all parameters.

wejax.jacobian.findiff(template, argnums, steps)[source]¶

Compute the Jacobian of a function using finite differences.

This function use numerical differentiation to compute the Jacobian of the template function,

\[J_{ij}(\theta)(f) = \frac{h_i(\theta + \delta_j)(f) - h_i(\theta - \delta_j)(f)}{2 \delta_j},\]

where \(h_i(\theta)\) is the template function, \(\theta\) is the model parameters, \(f\) is the frequency vector, and \(\delta_j\) is the step for numerical differentiation.

Example

>>> from wejax.jacobian import findiff
>>> jac = findiff(template, argnums=(0, 1), steps=(1e-6, 2e-6))
>>> jac(1.0, 2.0)
Array([[[[...]]]], dtype=complex64)

Parameters:

template (Callable[..., Array]) – Template function, implemented in JAX.
argnums (Sequence[int]) – Indices of the positional arguments to differentiate with respect to.
steps (TypeAliasType) –
Steps for numerical differentiation.

These can be either a scalar or an array with the same shape as the parameters to differentiate. If a scalar, the same step is used for all parameters.

Returns:

Jacobian function.

Return type:

Callable[..., Array]

Mixing Jacobians¶

The mix() function concatenates multiple Jacobians along the parameter axis, based on the given argument indices.

wejax.jacobian.mix(jacobians, argnums, sort_args=False)[source]¶

Concatenate Jacobians along the parameter axis.

This function concatenates multiple Jacobians along the parameter axis, based on the given argument indices.

Example

>>> from wejax.jacobian import autograd, findiff, mix
>>> jac1 = autograd(template1, argnums=(1,))
>>> jac2 = findiff(template2, argnums=(0,), steps=1e-6)
>>> jac = mix([jac1, jac2], argnums=[(1,), (0,)])
>>> associated_argnums = (0, 2, 1,)

You can sort the arguments before concatenating the Jacobians by setting sort_args=True.

>>> jac = mix([jac1, jac2], argnums=[(1,), (0,)], sort_args=True)
>>> associated_argnums = (0, 1)

Parameters:

jacobians (Sequence[Callable[..., Array]]) – List of Jacobian functions.
argnums (Sequence[Sequence[int]]) – List of lists of indices of argument indices attached to each Jacobian, in the same order as jacobians.
sort_args (bool, default: False) – Whether to sort the arguments (i.e., Jacobian columns) before concatenating the Jacobians.

Returns:

Mixed Jacobian function.

Return type:

Callable[..., Array]