First tutorial

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
from jax._src.typing import Array

jax.config.update("jax_enable_x64", True)

Setup

We imagine a simple template with 2 parameters $a$ and $b$, returning a single complex channel (depending on frequency) given by

$$ h(a, b)(f) = a^2 f + 42 i \cos(a) b^2 $$

The template function takes 2 parameters and returns an complex-valued array of shape (4, 1) for 4 frequency points.

For vectorization, one can use jax.vmap to apply the function element-wise to arrays of parameters. One needs to be careful to add dimensions on the left. As an example, one can use inputs of shape (N,) to get an output of shape (N, 4, 1) with

vtemplate = jax.vmap(template, in_axes=(0, 0), out_axes=0)

Therefore, the template output generalizes to (..., Nf, Nc) for $N_c$ channels, $N_f$ frequency points, and arbitrary additional dimensions from vectorization.

n_freqs = 4
n_channels = 1
n_params = 2

df = 0.01
freqs = np.arange(1, n_freqs + 1) * df


def template(a: float, b: float) -> Array:
    """Simple template function for testing purposes.

    Args:
        a: Parameter a.
        b: Parameter b.

    Returns:
        Single channel, shape (n_freqs, n_channels).
    """
    channel = a**2 * freqs + 42j * jnp.cos(a) * b**2
    return jnp.stack([channel], axis=-1)

def test_template_shapes() -> None:

    a = 1.0
    b = 1.0
    assert template(a, b).shape == (n_freqs, n_channels)

    a = jnp.asarray(1.0)
    b = jnp.asarray(1.0)
    assert template(a, b).shape == (n_freqs, n_channels)

    a = jnp.ones((50,))
    b = jnp.ones((50,))
    vtemplate = jax.vmap(template, in_axes=(0, 0))
    assert vtemplate(a, b).shape == (50, n_freqs, n_channels)

    a = jnp.ones((2, 3))
    b = jnp.ones((3,))
    a, b = jnp.broadcast_arrays(a, b)
    vtemplate = jax.vmap(template, in_axes=(0, 0))
    vvtemplate = jax.vmap(vtemplate, in_axes=(0, 0))
    assert vvtemplate(a, b).shape == (2, 3, n_freqs, n_channels)


test_template_shapes()

a = np.pi
b = 1 / np.sqrt(42)

template_val = template(a, b)

print("Template value:\n", template_val)
assert template_val.shape == (n_freqs, n_channels)
np.testing.assert_array_almost_equal(template_val.real[:, 0], np.pi**2 * freqs)
np.testing.assert_array_almost_equal(template_val.imag[:, 1], np.repeat(-1.0, n_freqs))

Template value:
 [[0.09869604-1.j]
 [0.19739209-1.j]
 [0.29608813-1.j]
 [0.39478418-1.j]]

Jacobian

We compute the analytic Jacobian with sympy.

That’s the analytic solution we will compare against.

sym_a, sym_b, sym_f = sp.symbols("a b f", real=True)
template_sym = sym_a**2 * sym_f + 42j * sp.cos(sym_a) * sym_b**2


jacobian_sym = sp.Matrix(
    [
        sp.diff(template_sym, sym_a),
        sp.diff(template_sym, sym_b),
    ]
)

jacobian_sym

\[\begin{split}\displaystyle \left[\begin{matrix}2 a f - 42.0 i b^{2} \sin{\left(a \right)}\\84.0 i b \cos{\left(a \right)}\end{matrix}\right]\end{split}\]

We expect the Jacobian of $h$ to be of shape (4, 1, 2) for 2 parameters (for which we take the partial derivatives), 1 channel, and 4 frequency points.

$$ \nabla h(a, b)(f)= \begin{bmatrix} 2 a f - 42i b^2 \sin(a) \ 84 i b \cos(a) \end{bmatrix} $$

Following the usual Numpy logic, we stack the partial derivatives on the last axis.

As a consequence, using JAX vectorization, this shape generalizes to (..., Nf, Nc, Np) for $N_p$ derivatives with respect to the parameters, $N_c$ channels, and $N_f$ frequency points,

vjac = jax.vmap(jac, in_axes=(0, 0), out_axes=0)

def jacobian_analytic(a: float, b: float) -> Array:
    """Analytic Jacobian of template function.

    Args:
        a: Parameter a.
        b: Parameter b.

    Returns:
        Jacobian, shape (n_freqs, n_channels, n_params).
    """
    # 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)


analytic_val = jacobian_analytic(a, b)

assert analytic_val.shape == (n_freqs, n_channels, n_params)

def test_jacobian_analytic_shapes() -> None:

    a = 1.0
    b = 1.0
    assert jacobian_analytic(a, b).shape == (n_freqs, n_channels, n_params)

    a = jnp.asarray(1.0)
    b = jnp.asarray(1.0)
    assert jacobian_analytic(a, b).shape == (n_freqs, n_channels, n_params)

    a = jnp.ones((50,))
    b = jnp.ones((50,))
    vjac = jax.vmap(jacobian_analytic, in_axes=(0, 0))
    assert vjac(a, b).shape == (50, n_freqs, n_channels, n_params)

    for i in range(50):
        np.testing.assert_array_almost_equal(
            vjac(a, b)[i], jacobian_analytic(a[i], b[i])
        )

    a = jnp.ones((2, 3))
    b = jnp.ones((3,))
    a, b = jnp.broadcast_arrays(a, b)
    vjac = jax.vmap(jacobian_analytic, in_axes=(0, 0))
    vvjac = jax.vmap(vjac, in_axes=(0, 0))
    assert vvjac(a, b).shape == (2, 3, n_freqs, n_channels, n_params)

    for i in range(2):
        for j in range(3):
            np.testing.assert_array_almost_equal(
                vvjac(a, b)[i, j], jacobian_analytic(a[i, j], b[i, j])
            )


test_jacobian_analytic_shapes()

We now use JAX’s autodiff to compute the Jacobian of $h$ at a given point $(a, b)$. We compare the result with the analytic solution.

from wejax.jacobian import autograd

jacobian_autograd = autograd(template, argnums=(0, 1))
autograd_val = jacobian_autograd(a, b)

assert autograd_val.shape == (n_freqs, n_channels, n_params)
print("Differences:\n", np.abs(autograd_val - analytic_val))
np.testing.assert_array_almost_equal(autograd_val, analytic_val)

Differences:
 [[[0. 0.]]

 [[0. 0.]]

 [[0. 0.]]

 [[0. 0.]]]

def test_jacobian_autograd_shapes() -> None:

    a = 1.0
    b = 1.0
    assert jacobian_autograd(a, b).shape == (n_freqs, n_channels, n_params)

    a = jnp.asarray(1.0)
    b = jnp.asarray(1.0)
    assert jacobian_autograd(a, b).shape == (n_freqs, n_channels, n_params)

    a = jnp.ones((50,))
    b = jnp.ones((50,))
    vjac = jax.vmap(jacobian_autograd, in_axes=(0, 0))
    assert vjac(a, b).shape == (50, n_freqs, n_channels, n_params)

    for i in range(50):
        np.testing.assert_array_almost_equal(
            vjac(a, b)[i], jacobian_autograd(a[i], b[i])
        )

    a = jnp.ones((2, 3))
    b = jnp.ones((3,))
    a, b = jnp.broadcast_arrays(a, b)
    vjac = jax.vmap(jacobian_autograd, in_axes=(0, 0))
    vvjac = jax.vmap(vjac, in_axes=(0, 0))
    assert vvjac(a, b).shape == (2, 3, n_freqs, n_channels, n_params)

    for i in range(2):
        for j in range(3):
            np.testing.assert_array_almost_equal(
                vvjac(a, b)[i, j], jacobian_autograd(a[i, j], b[i, j])
            )


test_jacobian_autograd_shapes()

We can also use numerical differentiation (finite differences) to check the correctness of the Jacobian.

from wejax.jacobian import findiff

jacobian_findiff = findiff(template, argnums=(0, 1), steps=[0.01, 0.1])
findiff_val = jacobian_findiff(a, b)

assert findiff_val.shape == (n_freqs, n_channels, n_params)
print("Differences:\n", np.abs(findiff_val - analytic_val))
np.testing.assert_array_almost_equal(findiff_val, analytic_val, decimal=5)

Differences:
 [[[2.11297566e-15 0.00000000e+00]]

 [[4.22062457e-15 0.00000000e+00]]

 [[8.38307840e-15 0.00000000e+00]]

 [[8.43858367e-15 0.00000000e+00]]]

def test_jacobian_findiff_shapes() -> None:

    a = 1.0
    b = 1.0
    assert jacobian_findiff(a, b).shape == (n_freqs, n_channels, n_params)

    a = jnp.asarray(1.0)
    b = jnp.asarray(1.0)
    assert jacobian_findiff(a, b).shape == (n_freqs, n_channels, n_params)

    a = jnp.ones((50,))
    b = jnp.ones((50,))
    vjac = jax.vmap(jacobian_findiff, in_axes=(0, 0))
    assert vjac(a, b).shape == (50, n_freqs, n_channels, n_params)

    for i in range(50):
        np.testing.assert_array_almost_equal(
            vjac(a, b)[i], jacobian_findiff(a[i], b[i])
        )

    a = jnp.ones((2, 3))
    b = jnp.ones((3,))
    a, b = jnp.broadcast_arrays(a, b)
    vjac = jax.vmap(jacobian_findiff, in_axes=(0, 0))
    vvjac = jax.vmap(vjac, in_axes=(0, 0))
    assert vvjac(a, b).shape == (2, 3, n_freqs, n_channels, n_params)

    for i in range(2):
        for j in range(3):
            np.testing.assert_array_almost_equal(
                vvjac(a, b)[i, j], jacobian_findiff(a[i, j], b[i, j])
            )


test_jacobian_findiff_shapes()

We can also built out our own Jacobian, mixing various methods.

from wejax.jacobian import mix

jacobian_mix = mix(
    [
        autograd(template, argnums=(1,)),
        findiff(template, argnums=(0,), steps=[0.01]),
    ],
    argnums=[(1,), (0,)],
    sort_args=True,
)

mix_val = jacobian_mix(a, b)

assert mix_val.shape == (n_freqs, n_channels, n_params)
print("Differences:\n", np.abs(mix_val - analytic_val))
np.testing.assert_array_almost_equal(mix_val, analytic_val, decimal=5)

Differences:
 [[[2.11297566e-15 0.00000000e+00]]

 [[4.22062457e-15 0.00000000e+00]]

 [[8.38307840e-15 0.00000000e+00]]

 [[8.43858367e-15 0.00000000e+00]]]

def test_jacobian_mix_shapes() -> None:

    a = 1.0
    b = 1.0
    assert jacobian_mix(a, b).shape == (n_freqs, n_channels, n_params)

    a = jnp.asarray(1.0)
    b = jnp.asarray(1.0)
    assert jacobian_mix(a, b).shape == (n_freqs, n_channels, n_params)

    a = jnp.ones((50,))
    b = jnp.ones((50,))
    vjac = jax.vmap(jacobian_mix, in_axes=(0, 0))
    assert vjac(a, b).shape == (50, n_freqs, n_channels, n_params)

    for i in range(50):
        np.testing.assert_array_almost_equal(vjac(a, b)[i], jacobian_mix(a[i], b[i]))

    a = jnp.ones((2, 3))
    b = jnp.ones((3,))
    a, b = jnp.broadcast_arrays(a, b)
    vjac = jax.vmap(jacobian_mix, in_axes=(0, 0))
    vvjac = jax.vmap(vjac, in_axes=(0, 0))
    assert vvjac(a, b).shape == (2, 3, n_freqs, n_channels, n_params)

    for i in range(2):
        for j in range(3):
            np.testing.assert_array_almost_equal(
                vvjac(a, b)[i, j], jacobian_mix(a[i, j], b[i, j])
            )


test_jacobian_mix_shapes()

FIM

We now have a Jacobian, which is a function of the 2 parameters $(a, b)$, and which returns an array of shape (4, 1, 2).

We saw that the (vectorized) Jacobian shapes generalizes to (..., Nf, Nc, Np) where $N_p$ is the number of parameters, $N_f$ the number of frequency points, and $N_c$ is the number of channels.

We want to build the FIM as

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

We first implement the inner product, contracting on the channel axis, and summing on the frequency axis.

from wejax.fim import inner_product


def test_inner_product_shape() -> None:

    a = jnp.ones((n_freqs, n_channels))
    b = jnp.ones((n_freqs, n_channels))
    noise_cov = jnp.ones((n_freqs, n_channels, n_channels))
    assert inner_product(a, b, noise_cov, df=df).shape == ()

    a = jnp.ones((n_freqs, n_channels))
    b = jnp.ones((n_freqs, n_channels))
    noise_cov = jnp.ones((n_channels, n_channels))
    assert inner_product(a, b, noise_cov, df=df).shape == ()

    a = jnp.ones((50, n_freqs, n_channels))
    b = jnp.ones((50, n_freqs, n_channels))
    noise_cov = jnp.ones((n_freqs, n_channels, n_channels))
    assert inner_product(a, b, noise_cov, df=df).shape == (50,)

    a = jnp.ones((20, n_freqs, n_channels))
    b = jnp.ones((n_freqs, n_channels))
    noise_cov = jnp.ones((n_freqs, n_channels, n_channels))
    assert inner_product(a, b, noise_cov, df=df).shape == (20,)


test_inner_product_shape()

We use vectorization to compute the FIM column by column, and stack the results.

from wejax.fim import matched_filter_gaussian_fim


def test_matched_filter_gaussian_fim_shape() -> None:

    a = 1.0
    b = 1.0
    noise_cov = jnp.ones((n_channels, n_channels))
    fim = matched_filter_gaussian_fim(jacobian_analytic, noise_cov, df=df)
    assert fim(a, b).shape == (n_params, n_params)

    a = 1.0
    b = 1.0
    noise_cov = jnp.ones((n_freqs, n_channels, n_channels))
    fim = matched_filter_gaussian_fim(jacobian_analytic, noise_cov, df=df)
    assert fim(a, b).shape == (n_params, n_params)

    a = jnp.ones((50,))
    b = jnp.ones((50,))
    noise_cov = jnp.ones((n_freqs, n_channels, n_channels))
    fim = matched_filter_gaussian_fim(jacobian_analytic, noise_cov, df=df)
    vfim = jax.vmap(fim, in_axes=(0, 0), out_axes=0)
    assert vfim(a, b).shape == (50, n_params, n_params)

    a = jnp.ones((20, 50))
    b = jnp.ones((50,))
    a, b = jnp.broadcast_arrays(a, b)
    noise_cov = jnp.ones((n_freqs, n_channels, n_channels))
    fim = matched_filter_gaussian_fim(jacobian_analytic, noise_cov, df=df)
    vfim = jax.vmap(fim, in_axes=(0, 0), out_axes=0)
    vvfim = jax.vmap(vfim, in_axes=(0, 0), out_axes=0)
    assert vvfim(a, b).shape == (20, 50, n_params, n_params)


test_matched_filter_gaussian_fim_shape()

We use Sympy to compute the analytic FIM.

fim_sym_00 = 4 * sp.re(
    sp.integrate(
        jacobian_sym[0] * sp.conjugate(jacobian_sym[0]), (sym_f, freqs[0], freqs[-1])
    )
)
fim_sym_01 = 4 * sp.re(
    sp.integrate(
        jacobian_sym[0] * sp.conjugate(jacobian_sym[1]), (sym_f, freqs[0], freqs[-1])
    )
)
fim_sym_10 = 4 * sp.re(
    sp.integrate(
        jacobian_sym[1] * sp.conjugate(jacobian_sym[0]), (sym_f, freqs[0], freqs[-1])
    )
)
fim_sym_11 = 4 * sp.re(
    sp.integrate(
        jacobian_sym[1] * sp.conjugate(jacobian_sym[1]), (sym_f, freqs[0], freqs[-1])
    )
)

assert fim_sym_01 == fim_sym_10

fim_sym = sp.Matrix([[fim_sym_00, fim_sym_01], [fim_sym_10, fim_sym_11]])
fim_sym

\[\begin{split}\displaystyle \left[\begin{matrix}0.000336 a^{2} + 211.68 b^{4} \sin^{2}{\left(a \right)} & - 423.36 b^{3} \sin{\left(a \right)} \cos{\left(a \right)}\\- 423.36 b^{3} \sin{\left(a \right)} \cos{\left(a \right)} & 846.72 b^{2} \cos^{2}{\left(a \right)}\end{matrix}\right]\end{split}\]

noise_cov = np.diag([1])

fim = matched_filter_gaussian_fim(jacobian_mix, noise_cov, df=df)
fim_val = fim(a, b)

assert fim_val.shape == (n_params, n_params)
print("FIM value:\n", fim_val)

expected_fim = sp.matrix2numpy(fim_sym.evalf(subs={sym_a: a, sym_b: b}), dtype=float)
assert expected_fim.shape == (n_params, n_params)
print("Expected FIM:\n", expected_fim)

print("Differences:\n", np.abs(fim_val - expected_fim))
np.testing.assert_array_almost_equal(fim_val, expected_fim, decimal=3)

FIM value:
 [[3.39514391e-03 0.00000000e+00]
 [0.00000000e+00 2.01600000e+01]]
Expected FIM:
 [[3.31618708e-03 1.90478840e-16]
 [1.90478840e-16 2.01600000e+01]]
Differences:
 [[7.89568352e-05 1.90478840e-16]
 [1.90478840e-16 0.00000000e+00]]

Errors

We can now compute the error covariance and correlation matrices.

from wejax.errors import covariance_matrix

cov = covariance_matrix(fim=fim_val, argnums=(0, 1))

def test_covariance_matrix_shape() -> None:

    a = 1.0
    b = 1.0
    noise_cov = np.ones((n_channels, n_channels))
    a = 1.0
    b = 1.0
    noise_cov = np.ones((n_channels, n_channels))
    fim = matched_filter_gaussian_fim(jacobian_analytic, noise_cov, df=df)
    assert covariance_matrix(fim=fim(a, b)).shape == (n_params, n_params)
    assert covariance_matrix(fim=fim(a, b)).shape == (n_params, n_params)
    assert covariance_matrix(fim=fim(a, b), argnums=(0,)).shape == (1, 1)
    assert covariance_matrix(fim=fim(a, b), argnums=(1,)).shape == (1, 1)
    assert covariance_matrix(fim=fim(a, b), argnums=(0, 1, 0)).shape == (3, 3)

    a = 1.0
    b = 1.0
    noise_cov = np.ones((n_freqs, n_channels, n_channels))
    a = 1.0
    b = 1.0
    noise_cov = np.ones((n_freqs, n_channels, n_channels))
    fim = matched_filter_gaussian_fim(jacobian_analytic, noise_cov, df=df)
    assert covariance_matrix(fim=fim(a, b)).shape == (n_params, n_params)
    assert covariance_matrix(fim=fim(a, b)).shape == (n_params, n_params)

    a = np.ones((50,))
    b = np.ones((50,))
    noise_cov = np.ones((n_freqs, n_channels, n_channels))
    vjac = jax.vmap(jacobian_analytic, in_axes=(0, 0))
    fim = matched_filter_gaussian_fim(vjac, noise_cov, df=df)
    assert covariance_matrix(fim=fim(a, b)).shape == (50, n_params, n_params)


test_covariance_matrix_shape()

from wejax.errors import stddevs

stddev_from_fim = stddevs(fim=fim_val)
stddev_from_cov = stddevs(cov=cov)
expected = np.sqrt(np.diag(np.linalg.inv(fim_val)))

np.testing.assert_array_almost_equal(stddev_from_fim, expected)
np.testing.assert_array_almost_equal(stddev_from_cov, expected)

def test_stddevs_shape() -> None:

    a = 1.0
    b = 1.0
    noise_cov = np.ones((n_channels, n_channels))
    fim = matched_filter_gaussian_fim(jacobian_analytic, noise_cov, df=df)
    assert stddevs(fim=fim(a, b)).shape == (n_params,)

    a = np.ones((50,))
    b = np.ones((50,))
    noise_cov = np.ones((n_freqs, n_channels, n_channels))
    vjac = jax.vmap(jacobian_analytic, in_axes=(0, 0))
    fim = matched_filter_gaussian_fim(vjac, noise_cov, df=df)
    assert stddevs(fim=fim(a, b)).shape == (50, n_params)


test_stddevs_shape()

from wejax.errors import correlation_matrix

corr_from_fim = correlation_matrix(fim=fim_val)
corr_from_cov = correlation_matrix(cov=cov)

np.testing.assert_array_almost_equal(corr_from_fim, corr_from_cov)
np.testing.assert_array_almost_equal(np.diag(corr_from_fim), stddevs(fim=fim_val))

def test_correlation_matrix_shape() -> None:

    a = 1.0
    b = 1.0
    noise_cov = np.ones((n_channels, n_channels))
    fim = matched_filter_gaussian_fim(jacobian_analytic, noise_cov, df=df)
    assert correlation_matrix(fim=fim(a, b)).shape == (n_params, n_params)

    a = np.ones((50,))
    b = np.ones((50,))
    noise_cov = np.ones((n_freqs, n_channels, n_channels))
    vjac = jax.vmap(jacobian_analytic, in_axes=(0, 0))
    fim = matched_filter_gaussian_fim(vjac, noise_cov, df=df)
    assert correlation_matrix(fim=fim(a, b)).shape == (50, n_params, n_params)


test_correlation_matrix_shape()

from wejax.errors import plot_correlation_matrix

corr = np.array([[4.23e-1, 0.5, 0.2], [0.5, 3e-6, -0.7], [0.2, -0.7, 3.2e2]])

plot_correlation_matrix(
    corr,
    title="Test correlation matrix",
    param_names=["$f_0$", "$d_L$", "$\phi_0$"],
    param_units=["Hz", "pc", "rad"],
)

(<Figure size 640x480 with 2 Axes>,
 <Axes: title={'center': 'Test correlation matrix'}>)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))

n_params = 6
corr = np.random.uniform(-1, 1, (n_params, n_params))
corr += np.diag(np.exp(np.random.uniform(1, 10, n_params)))

plot_correlation_matrix(corr, ax=ax1, title="Source 1", show_cbar=False)

n_params = 10
corr = np.random.uniform(-1, 1, (n_params, n_params))
corr += np.diag(np.exp(np.random.uniform(1, 10, n_params)))
names = [f"Param {i}" for i in range(n_params)]
units = [f"Unit {i}" for i in range(n_params)]

plot_correlation_matrix(
    corr,
    ax=ax2,
    title="Source 2",
    font_size=6,
    show_cbar=False,
    param_names=names,
    label_rotation=45,
)

(<Figure size 1200x600 with 2 Axes>, <Axes: title={'center': 'Source 2'}>)