Skip to content

metrics

Differentiable loss functions for parameter fitting and model comparison.

All functions in this module operate on plain jax.Array arguments and are fully differentiable through jax.grad. They are designed to be passed as (or called within) the loss_fn argument of fit_parameters.

from jbubble.metrics import mse_radius, normalised_mse_radius
from jbubble.metrics import mse_emission, normalised_mse_emission
from jbubble.metrics import peak_expansion, peak_expansion_error

mse_radius 

mse_radius(r_sim: Array, r_target: Array) -> Array

Mean squared error between simulated and target radii [m²].

\[ \frac{1}{N}\sum_{i=1}^{N} \left(R_{\text{sim},i} - R_{\text{target},i}\right)^2 \]

Parameters:

Name Type Description Default
r_sim (Array, shape(N))

Simulated radius trajectory [m].

 required
r_target (Array, shape(N))

Target radius trajectory [m].

 required

Returns:

Type Description
Array

Scalar mean squared error [m²].
Source code in jbubble/metrics.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def mse_radius(r_sim: jax.Array, r_target: jax.Array) -> jax.Array:
    r"""Mean squared error between simulated and target radii [m²].

    $$
    \frac{1}{N}\sum_{i=1}^{N} \left(R_{\text{sim},i} - R_{\text{target},i}\right)^2
    $$

    Parameters
    ----------
    r_sim : jax.Array, shape (N,)
        Simulated radius trajectory [m].
    r_target : jax.Array, shape (N,)
        Target radius trajectory [m].

    Returns
    -------
    jax.Array
        Scalar mean squared error [m²].
    """
    return jnp.mean((r_sim - r_target) ** 2)

normalised_mse_radius 

normalised_mse_radius(r_sim: Array, r_target: Array, R0: ArrayLike) -> Array

Normalised mean squared radius error (dimensionless).

\[ \frac{1}{N}\sum_{i=1}^{N} \left(\frac{R_{\text{sim},i} - R_{\text{target},i}}{R_0}\right)^2 \]

Equivalent to mse_radius(r_sim / R0, r_target / R0). The \(R_0\) normalisation makes the loss dimensionless and of order 1 regardless of bubble size, which improves optimiser conditioning.

Parameters:

Name Type Description Default
r_sim (Array, shape(N))

Simulated radius trajectory [m].

 required
r_target (Array, shape(N))

Target radius trajectory [m].

 required
R0 float or Array

Equilibrium radius used for normalisation [m].

 required

Returns:

Type Description
Array

Scalar normalised mean squared error.
Source code in jbubble/metrics.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def normalised_mse_radius(
    r_sim: jax.Array,
    r_target: jax.Array,
    R0: ArrayLike,
) -> jax.Array:
    r"""Normalised mean squared radius error (dimensionless).

    $$
    \frac{1}{N}\sum_{i=1}^{N}
        \left(\frac{R_{\text{sim},i} - R_{\text{target},i}}{R_0}\right)^2
    $$

    Equivalent to `mse_radius(r_sim / R0, r_target / R0)`. The $R_0$
    normalisation makes the loss dimensionless and of order 1 regardless of
    bubble size, which improves optimiser conditioning.

    Parameters
    ----------
    r_sim : jax.Array, shape (N,)
        Simulated radius trajectory [m].
    r_target : jax.Array, shape (N,)
        Target radius trajectory [m].
    R0 : float or jax.Array
        Equilibrium radius used for normalisation [m].

    Returns
    -------
    jax.Array
        Scalar normalised mean squared error.
    """
    return jnp.mean(((r_sim - r_target) / R0) ** 2)

peak_expansion 

peak_expansion(r_sim: Array, R0: ArrayLike) -> Array

Maximum radial expansion ratio, \(R_\text{max}/R_0\).

Parameters:

Name Type Description Default
r_sim (Array, shape(N))

Simulated radius trajectory [m].

 required
R0 float or Array

Equilibrium radius [m].

 required

Returns:

Type Description
Array

Scalar expansion ratio (dimensionless).
Source code in jbubble/metrics.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def peak_expansion(r_sim: jax.Array, R0: ArrayLike) -> jax.Array:
    r"""Maximum radial expansion ratio, $R_\text{max}/R_0$.

    Parameters
    ----------
    r_sim : jax.Array, shape (N,)
        Simulated radius trajectory [m].
    R0 : float or jax.Array
        Equilibrium radius [m].

    Returns
    -------
    jax.Array
        Scalar expansion ratio (dimensionless).
    """
    return jnp.max(r_sim) / R0

peak_expansion_error 

peak_expansion_error(
    r_sim: Array, R0: ArrayLike, target_expansion: ArrayLike
) -> Array

Squared error in the peak expansion ratio.

\[ \left(\frac{R_\text{max}}{R_0} - \mathrm{target\_expansion}\right)^2 \]

Useful when only the maximum oscillation amplitude matters rather than the full waveform.

Parameters:

Name Type Description Default
r_sim (Array, shape(N))

Simulated radius trajectory [m].

 required
R0 float or Array

Equilibrium radius [m].

 required
target_expansion float or Array

Target \(R_\text{max}/R_0\) value.

 required

Returns:

Type Description
Array

Scalar squared error (dimensionless).
Source code in jbubble/metrics.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def peak_expansion_error(
    r_sim: jax.Array,
    R0: ArrayLike,
    target_expansion: ArrayLike,
) -> jax.Array:
    r"""Squared error in the peak expansion ratio.

    $$
    \left(\frac{R_\text{max}}{R_0} - \mathrm{target\_expansion}\right)^2
    $$

    Useful when only the maximum oscillation amplitude matters rather
    than the full waveform.

    Parameters
    ----------
    r_sim : jax.Array, shape (N,)
        Simulated radius trajectory [m].
    R0 : float or jax.Array
        Equilibrium radius [m].
    target_expansion : float or jax.Array
        Target $R_\text{max}/R_0$ value.

    Returns
    -------
    jax.Array
        Scalar squared error (dimensionless).
    """
    return (peak_expansion(r_sim, R0) - target_expansion) ** 2

mse_emission 

mse_emission(p_sim: Array, p_target: Array) -> Array

Mean squared error between simulated and target emission pressure [Pa²].

Parameters:

Name Type Description Default
p_sim (Array, shape(N))

Simulated radiated pressure [Pa].

 required
p_target (Array, shape(N))

Target radiated pressure [Pa].

 required

Returns:

Type Description
Array

Scalar mean squared error [Pa²].
Source code in jbubble/metrics.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def mse_emission(
    p_sim: jax.Array,
    p_target: jax.Array,
) -> jax.Array:
    """Mean squared error between simulated and target emission pressure [Pa²].

    Parameters
    ----------
    p_sim : jax.Array, shape (N,)
        Simulated radiated pressure [Pa].
    p_target : jax.Array, shape (N,)
        Target radiated pressure [Pa].

    Returns
    -------
    jax.Array
        Scalar mean squared error [Pa²].
    """
    return jnp.mean((p_sim - p_target) ** 2)

normalised_mse_emission 

normalised_mse_emission(
    p_sim: Array, p_target: Array, p_ref: ArrayLike
) -> Array

Normalised mean squared error (MSE) for emission pressure (dimensionless).

\[ \frac{1}{N}\sum_{i=1}^{N} \left(\frac{p_{\text{sim},i} - p_{\text{target},i}}{p_\text{ref}}\right)^2 \]

Parameters:

Name Type Description Default
p_sim (Array, shape(N))

Simulated radiated pressure [Pa].

 required
p_target (Array, shape(N))

Target radiated pressure [Pa].

 required
p_ref float or Array

Reference pressure for normalisation [Pa].

 required

Returns:

Type Description
Array

Scalar normalised mean squared error.
Source code in jbubble/metrics.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
def normalised_mse_emission(
    p_sim: jax.Array,
    p_target: jax.Array,
    p_ref: ArrayLike,
) -> jax.Array:
    r"""Normalised mean squared error (MSE) for emission pressure (dimensionless).

    $$
    \frac{1}{N}\sum_{i=1}^{N}
        \left(\frac{p_{\text{sim},i} - p_{\text{target},i}}{p_\text{ref}}\right)^2
    $$

    Parameters
    ----------
    p_sim : jax.Array, shape (N,)
        Simulated radiated pressure [Pa].
    p_target : jax.Array, shape (N,)
        Target radiated pressure [Pa].
    p_ref : float or jax.Array
        Reference pressure for normalisation [Pa].

    Returns
    -------
    jax.Array
        Scalar normalised mean squared error.
    """
    return jnp.mean(((p_sim - p_target) / p_ref) ** 2)

Choosing a loss function

Loss Use when
mse_radius Direct radius fitting; same units as data
normalised_mse_radius Fitting across bubbles of different sizes; scale-invariant
peak_expansion_error Only the amplitude matters, not the full waveform shape
mse_emission Fitting to hydrophone measurements
normalised_mse_emission Fitting to hydrophone data with a known reference level

Normalised losses are generally preferable for gradient-based fitting because they keep the loss \(\mathcal{O}(1)\) regardless of bubble size or driving pressure, making the choice of learning rate more robust.

Custom loss functions

Any differentiable function of SimulationResult fields works as a loss_fn. For example, fitting to the power spectral density:

import jax.numpy as jnp


def psd_loss(result, target_psd):
    fft_r = jnp.abs(jnp.fft.rfft(result.radius)) ** 2
    return jnp.mean((fft_r - target_psd) ** 2)


fit_result = fit_parameters(
    ...,
    loss_fn=lambda result: psd_loss(result, measured_psd),
)