Fisher information matrix

Use this module to compute Fisher Information Matrices (FIM) for set of observables, assuming a Gaussian-noise matched filter-like likelihood.

Noise-weighted inner products

Noise-weighted inner products are used to compute the Fisher Information Matrix. They are defined as

\[(a|b) = 4 \Re \int_{f_{\min}}^{f_{\max}} \tilde{a}(f) \Sigma^{-1}(f) \tilde{b}^+(f) \, \mathrm{d}f,\]

where \(\tilde{a}(f)\) and \(\tilde{b}(f)\) are single or multi-channel complex-valued Fourier transforms of the observables \(a(t)\) and \(b(t)\). The noise covariance matrix \(\Sigma(f)\) characterizes the noise correlations between the various observable channels.

wejax.fim.univariate_inner_product(a, b, noise_psd, *, df, sum_func=<PjitFunction of <function trapezoid>>)

Compute the noise-weighted inner product of two univariate observables.

The inner product is computed in the frequency domain, and for two univariate real-valued observables. The inner product is defined as

\[(a|b) = 4 \Re \int_0^\infty \frac{\tilde{a}(f) \tilde{b}^*(f)}{S(f)} \, \mathrm{d}f,\]

where \(\tilde{a}(f)\) and \(\tilde{b}(f)\) are the Fourier transforms of the observables \(a(t)\) and \(b(t)\), respectively, and \(S(f)\) is the one-sided power spectral density of the noise.

Note that we assume here that the noise is stationary, such that there is no correlation between the noise at different frequencies; ie., the noise power spectral density characterizes the noise power at each frequency.

For multivariate observables, use inner_product() instead.

Example

>>> result = univariate_inner_product(a, b, my_psd, df=df)

Parameters:

• a (Array) –

  Discrete Fourier transforms for the observables to compute the inner product of. Since the signals are assumed to be real, we only need the positive frequencies, as returned by jax.numpy.fft.rfft().

  Here, Nf is the number of frequencies and ... are any additional dimensions.

• b (Array) –

  Discrete Fourier transforms for the observables to compute the inner product of. Since the signals are assumed to be real, we only need the positive frequencies, as returned by jax.numpy.fft.rfft().

  Here, Nf is the number of frequencies and ... are any additional dimensions.

• noise_psd (Array) –

  Noise covariance matrix [/Hz].

  Here Nf is the number of frequencies and ... are any additional dimensions.

• df (float) – Frequency spacing [Hz].

• sum_func (Callable, default: <PjitFunction of <function trapezoid at 0x7ca1afd9fa60>>) –

  Function to use to sum over the frequencies.

  The function should take as first positional argument an JAX Array, and should accept the keyword arguments dx and axis. It should return an array of shape (...), where ... are any additional dimensions.

  We suggest using jax.numpy.trapezoid(), or jax.numpy.sum().

Returns:

Noise-weighted inner product.

Return type:

Array

wejax.fim.inner_product(a, b, noise_cov, *, df, should_invert_noise_cov=True, sum_func=<PjitFunction of <function trapezoid>>)

Compute the noise-weighted inner product of two sets of observables.

The inner product is computed in the frequency domain, and for two sets of real-valued observables. The inner product is defined as

\[(a|b) = 4 \Re \int_0^\infty \tilde{a}(f) \Sigma^{-1}(f) \tilde{b}^+(f) \, \mathrm{d}f,\]

where \(\tilde{a}(f)\) and \(\tilde{b}(f)\) are the Fourier transforms of the sets of observables \(a(t)\) and \(b(t)\), respectively, and \(\Sigma(f)\) is the noise covariance matrix.

By default, the inverse of the noise covariance matrix is computed and used to compute the inner product. If you want to provide the inverse noise covariance matrix directly, set should_invert_noise_cov to False.

Note that we assume here that the noise is stationary, such that there is no correlation between the noise at different frequencies (i.e., the noise covariance matrix is frequency-dependent and only characterize the noise correlations between the various observables channels).

For univariate observables, use univariate_inner_product() instead.

Example

>>> result = inner_product(a, b, noise_cov=noise_cov, df=df)

Parameters:

• a (Array) –

  Discrete Fourier transforms for the observables to compute the inner product of. Since the signals are assumed to be real, we only need the positive frequencies, as returned by jax.numpy.fft.rfft().

  Here, Nf is the number of frequencies, Nc is the number of observable channels, and ... are any additional dimensions.

• b (Array) –

  Discrete Fourier transforms for the observables to compute the inner product of. Since the signals are assumed to be real, we only need the positive frequencies, as returned by jax.numpy.fft.rfft().

  Here, Nf is the number of frequencies, Nc is the number of observable channels, and ... are any additional dimensions.

• noise_cov (Array) –

  Noise covariance matrix [/Hz].

  Here Nf is the number of frequencies, Nc is the number of observable channels, and ... are any additional dimensions.

• df (float) – Frequency spacing [Hz].

• should_invert_noise_cov (bool, default: True) – If True, the inverse of the noise covariance matrix is computed and used to compute the inner product. If False, the provided noise_cov is used directly as the inverse noise covariance matrix.

• sum_func (Callable, default: <PjitFunction of <function trapezoid at 0x7ca1afd9fa60>>) –

  Function to use to sum over the frequencies.

  The function should take as first positional argument an JAX Array, and should accept the keyword arguments dx and axis. It should return an array of shape (...), where ... are any additional dimensions.

  We suggest using jax.numpy.trapezoid(), or jax.numpy.sum().

Returns:

Noise-weighted inner product.

Return type:

Array

FIM computation

Use the matched_filter_gaussian_fim() function to compute the Fisher Information Matrix from a Jacobian function. The function returns another function, which can evaluate the Fisher Information Matrix for a set of parameters.

wejax.fim.matched_filter_gaussian_fim(jacobian, noise_cov, *, df, should_invert_noise_cov=True, sum_func=<PjitFunction of <function trapezoid>>)

Return a function that computes the Fisher Information Matrix.

Compute the Fisher Information Matrix for a template function, assuming a stationary Gaussian noise. The Fisher Information Matrix is given by

\[I_{ij}(\theta) = \left( \frac{\partial h}{\partial \theta_i} \middle| \frac{\partial h}{\partial \theta_j} \right).\]

The returned function can be called with a set of parameters, and will return the Fisher Information Matrix evaluated at those parameters.

Example

>>> fim = matched_filter_gaussian_fim(jacobian, noise_cov=noise_cov, df=df)
>>> result_1 = fim(1.0, 2.0)
>>> result_2 = fim(3.0, 4.0)

The returned FIM function can be vectorized for efficient computation of multiple sets of parameters using jax.vmap() (see JAX documentation).

>>> vmap_fim = jax.vmap(fim, in_axes=(0, 0), out_axes=0)

Parameters:

• jacobian (Callable[..., Array]) – Jacobian function.

• noise_cov (Array) –

  Noise covariance matrix [/Hz].

  Here Nf is the number of frequencies, Nc is the number of observable channels, and ... are any additional dimensions.

• df (float) – Frequency spacing [Hz].

• should_invert_noise_cov (bool, default: True) – If True, the inverse of the noise covariance matrix is computed and used to compute the inner product. If False, the provided noise_cov is used directly as the inverse noise covariance matrix.

• sum_func (Callable, default: <PjitFunction of <function trapezoid at 0x7ca1afd9fa60>>) –

  Function to use to sum over the frequencies.

  The function should take as first positional argument an JAX Array, and should accept the keyword arguments dx and axis. It should return an array of shape (...), where ... are any additional dimensions.

  We suggest using jax.numpy.trapezoid(), or jax.numpy.sum().

Returns:

A function that computes the Fisher Information Matrix for a set of parameters.

Return type:

FIMFunction

class wejax.fim.FIMFunction(*args, **kwargs)

Protocol for Fisher Information Matrices functions.

Parameters:

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

Returns:

The Fisher Information Matrix evaluated at args.

Here Np is the number of parameters, and ... are any additional dimensions returned by the Jacobian function.

Return type:

Array of shape (..., Np, Np)