diabayes package#

Subpackages#

diabayes.gui package

Submodules#

diabayes.SVI module#

diabayes.SVI.compute_phi(x: Float[Array, 'N M'], gradp: Float[Array, 'N M'], gradq: Float[Array, 'N M']) → Float[Array, 'N M'][source]#

Compute the Stein variational gradients for a set of particles (x) and the gradients of the log-likelihood function (gradp) and the log-prior (gradq).

Parameters:

x (ParticleArray) – The set of invertible parameters (“particles”) of shape (Nparticles, Ndimensions)
gradp (GradientArray) – The gradients of the log-likelihood (gradp) and the log-prior (gradq) with respect to the invertible parameters. Has a shape (Nparticles, Ndimensions)
gradq (GradientArray) – The gradients of the log-likelihood (gradp) and the log-prior (gradq) with respect to the invertible parameters. Has a shape (Nparticles, Ndimensions)

Returns:

grad_x – The directional gradients for the particle updates, e.g. new_x = x + step_size * grad_x. Has the same shape as x and gradp, gradq.

Return type:

Gradients

diabayes.SVI.mapped_log_likelihood(log_params: RSFParams | CNSParams, mu_obs: Float[Array, 'Nt'], noise_std: Float, forward_fn: Callable[[RSFParams | CNSParams], Variables]) → Float#: Vectorized version of _log_likelihood. Takes similar arguments as _log_likelihood but with additional array axes over which _log_likelihood is mapped.

diabayes.forward_models module#

class diabayes.forward_models.Forward(friction_model: FrictionModel, state_evolution: Dict[str, Callable], stress_transfer: StressTransfer[BC], sundman: Callable | None = None)[source]#

Bases: Generic[BC]

The Forward class assembles the various components that comprise a forward model, such that Forward.__call__ takes some variables and returns the rate of change of these variables

\[\frac{\mathrm{d} \vec{X}}{\mathrm{d}t} = f \left( \vec{X} \right)\]

This forward ODE can then be solved by any ODE solver.

The Forward class is instantiated by providing a friction model, a “state” evolution law, and a stress transfer model.

Examples

>>> from diabayes.forward_models import ageing_law, rsf, springblock, Forward
>>> foward_model = Forward(rsf, {"theta": ageing_law}, springblock)
>>> X_dot = forward_model(variables=..., params=..., friction_constants=..., block_constants=...)

call(t: Float, variables: Variables, params: RSFParams | CNSParams, friction_constants: RSFConstants | CNSConstants, block_constants: BC) → Variables#

Calculate the rate of change of the variables, defining the ODE to be solved.

Parameters:

t (Float) – Current value of time [s]
variables (diabayes.Variables) – The instantaneous values of the variables for which the time derivative will be computed
params (_Params) – The forward model (invertible) parameters
friction_constants (_Constants) – The forward model constants
block_constants (_BlockConstants) – The constants associated with the stress transfer

Returns:

dvars – The instantaneous rate of change of the variables

Return type:

Variables

set_initial_values(**kwargs) → None[source]#

Set the initial values of the variables. This sets the values as self.variables, which is a Variables object.

Parameters:: **kwargs (dict) – Key-value pairs of variable names and corresponding values

diabayes.solver module#

class diabayes.solver.ODESolver(forward_model: Forward, rtol: float = 1e-07, atol: float = 1e-10, checkpoints: int = 100)[source]#

Bases: object

The main solver class that contains forward and inverse modelling methods.

bayesian_inversion(t: Float[Array, 'Nt'], mu: Float[Array, 'Nt'], noise_std: Float, y0: Variables, params: RSFParams | CNSParams, friction_constants: RSFConstants | CNSConstants, block_constants: SpringBlockConstants | InertialSpringBlockConstants, Nparticles: int = 1000, Nsteps: int = 150, rng: None | int | Array = None) → BayesianSolution[source]#

A Bayesian inversion routine using the Stein Variational Inference method.

Parameters:

t (Array) – A vector or time values (in units of seconds). The time steps do not need to be uniform
mu (Array) – The observed friction curve sampled at t
noise_std (Float) – An estimate of the standard deviation of the noise in the measured friction curve. A conservative value is recommended, i.e. if the noise has a standard deviation of $10^{-3}$, a good starting point would be to set noise_std = 0.5e-3
y0 (Variables) – The initial values for the modelled friction and any state variables
params (_Params) – The initial guess for the invertible parameters that characterise the forward problem. It is recommended to use the result from ODESolver.max_likelihood_inversion. The prior distribution will be centered around this initial guess
friction_constants (_Constants) – The non-invertible constants that characterise the forward problem
block_constants (_BlockConstants) – The stress transfer constants (e.g. stiffness and loading rate)
Nparticles (int) – The number of particles (= posterior samples) to include. A higher value gives a more accurate estimation of the posterior distribution, at a higher computational cost.
Nsteps (int) – The number of iterations before convergence is expected to be achieved. It is recommended to start with a value of 100 and then see if an equilibrium was actually achieved. Increasing this value beyond the point of equilibrium does not do anything.
rng (None, int, jax.random.PRNGKey) – The random seed used to initialise the particle swarm distribution. If None, the current time will be used as a seed, which leads to different result for each realisation. When an integer is provided, a new jax.random.PRNGKey is generated.

Returns:

result – The result of the inversion incapsulated in a BayesianSolution container, which provides access to the full convergence chains, as well as diagnostic and visualisation routines.

Return type:

diabayes.BayesianSolution

Notes

This method is currently only compatible with standard rate-and-state friction models.

evaluate_at_t(sol: Solution, t_eval: Float[Array, 'Nt']) → Variables[source]#

generate_sequence(t_steps: Float[Array, 'Nsteps'], v_steps: Float[Array, 'Nsteps'], dt: Float, y0: Variables, params: RSFParams | CNSParams, friction_constants: RSFConstants | CNSConstants, block_constants: SpringBlockConstants | InertialSpringBlockConstants, method: str = 'RK45') → Tuple[Variables, Float[Array, 'Nt']][source]#

Solve the forward problem for a sequence of velocity steps. For each load-point velocity in v_steps, a forward simulation is run using the previous step’s final state as the initial state.

This routine can be used to generate a sequence of (up/down) velocity steps, or a slide-hold-slide sequence (by setting a given v_step to zero).

Parameters:

t_steps (Float[Array, "Nsteps"]) – A vector of change point times of each step in the sequence
v_steps (Float[Array, "Nsteps"]) – A vector of the load-point velocity values for each step in the sequence
dt (Float) – The desired time sample spacing of the solution
y0 (Variables) – The initial values (fricton and state) wrapped in a Variables container.
params (_Params) – The (invertible) parameters that govern the dynamics, wrapped in a Params container.
friction_constants (_Constants) – A container object containing the friction constants
block_constants (_BlockConstants) – A container object containing the block constants. Note that the load-point velocity will be updated for each step

Returns:

result (Variables) – Solution time series of friction and state
t (Float[Array, “Nt”]) – Solution time samples

max_likelihood_inversion(t: Float[Array, 'Nt'], mu: Float[Array, 'Nt'], y0: Variables, params: RSFParams | CNSParams, friction_constants: RSFConstants | CNSConstants, block_constants: SpringBlockConstants | InertialSpringBlockConstants, verbose: bool = False, retries: int = 3, seed: int = 42) → Solution | None[source]#

Minimises the least-squares residuals between the observed friction curve and the parameterised one, using the Levenberg-Marquardt algorithm.

Parameters:

t (Array) – A vector or time values (in units of seconds). The time steps do not need to be uniform
mu (Array) – The observed friction curve sampled at t
y0 (Variables) – The initial values for the modelled friction and any state variables
params (_Params) – The initial guess for the invertible parameters that characterise the forward problem. These need to be sufficiently close to the “true” values for the algorithm to converge
friction_constants (_Constants) – The non-invertible constants that characterise the forward problem
block_constants (_BlockConstants) – The stress transfer constants (e.g. stiffness and loading rate)
verbose (bool) – Whether or not to output detailed progress of the inversion. Defaults to False
retries (int) – The maximum number of inversion attempts. When the inversion fails to converge, it will retry up to retries times with randomly perturbed initial parameters.
seed (int) – Seed for the random number generator. This is only used when the initial inversion attempt fails, and the initial parameters are randomly perturbed before the next attempt.

Returns:

sol – The inversion result, including various diagnostics. The inverted parameter values can be accessed as sol.values

Return type:

optimistix.Solution

solve_forward(t: Float[Array, 'Nt'], y0: Variables, params: RSFParams | CNSParams, friction_constants: RSFConstants | CNSConstants, block_constants: SpringBlockConstants | InertialSpringBlockConstants, method: str = 'RK45', interpolate_time: bool = True) → Variables[source]#

Solve a forward problem using SciPy’s solve_ivp routine. While this routine doesn’t propagate any gradients, it is much faster to initialise and to perform a single forward run. Hence for playing around with different parameters, it is preferred over a JITed JAX implementation.

Parameters:

t (Float[Array, "Nt"]) – A vector of time samples where a solution is requested.
y0 (Variables) – The initial values (fricton and state) wrapped in a Variables container.
params (_Params) – The (invertible) parameters that govern the dynamics, wrapped in a Params container.
friction_constants (_Constants) – A container object containing the friction constants
block_constants (_BlockConstants) – A container object containing the block constants
method (str) – The solver used by SciPy’s solve_ivp. Default: RK45
interpolate_time (bool) – Whether to interpolate the result to the user-provided time samples (True), or to use the samples from the adaptive ODE solver (False). Default: True

Returns:

result – Solution time series of friction and state

Return type:

Variables

atol: float#: The absolute tolerance used by the ODE solver

checkpoints: int#: The number of checkpoints to use to compute the (adjoint) gradients through the ODE routine. A higher number increases stability and speed, at the expense of more GPU memory

forward_model: Forward#: An instantiated diabayes.Forward class, including friction law, state evolution equations, and stress transfer equation

learning_rate: float = 0.01#: The initial learning rate provided to the Adam algorithm for the Stein Variational Inference. The default value of 1e-2 seems like a sensible choice for most models.

rtol: float#: The absolute tolerance used by the ODE solver

diabayes.typedefs module#

class diabayes.typedefs.BC#: Function signatures#

alias of TypeVar(‘BC’, bound=SpringBlockConstants | InertialSpringBlockConstants, contravariant=True)

class diabayes.typedefs.BayesianSolution(log_params: RSFParams | CNSParams, log_likelihood: Float[Array, 'Nsteps'], nan_count: Float[Array, 'Nsteps'])[source]#

Bases: object

The result of a Bayesian inversion. This data class stores the particle trajectories (equivalent of MCMC chains), the final equilibrum state, and various statistcs, as well as helper routines to diagnose or visualise the results.

Examples

>>> result = solver.bayesian_inversion(...)
>>> print(result.nan_count)  # Check for NaN values
>>> print(result.statistics.Dc.median)  # Get the median of Dc
>>> result.plot_convergence()  # Visually inspect convergence
>>> result.cornerplot(nbins=15)  # Create a cornerplot

cornerplot(nbins=20)[source]#

Create a corner plot from the particle positions. For N parameters, a corner plot is the lower half of an N x N grid of subplots. On the diagonal, the histogram of the marginal posterior over the i-th parameter is plotted. Below the diagonal, the (i,j)-th panel is a scatter plot of parameter i against parameter j, which shows the co-variance between the two quantities.

Parameters:: nbins (int) – The number of bins to use for the histograms on the diagonal of the corner plot
Returns:: fig – The figure object that can be manipulated after creation (e.g. to save to disk)
Return type:: matplotlib.pyplot.figure

plot_convergence()[source]#

Helper routine to visualise the convergence of the inversion. Convergence is achieved when the particles settle in an equilibrium position (i.e., they stop moving).

Returns:: fig – The figure object that can be manipulated after creation (e.g. to save to disk)
Return type:: matplotlib.pyplot.figure

sample(solver, y0: Variables, t: Float[Array, 'Nt'], friction_constants: RSFConstants | CNSConstants, block_constants: SpringBlockConstants | InertialSpringBlockConstants, nsamples: int = 10, rng: None | int | Array = None)[source]#

Draw random samples from the posterior distribution. Each particle should represent one realisation of a plausible friction curve. For nsamples realisations requested by a user, it is more beneficial to vmap the forward simulations rather than compute them one by one.

The resulting friction curves can be used for validation purposes (does the posterior match the data?).

Examples

>>> result = solver.bayesian_inversion(...)
>>> samples, sample_results = result.sample(
...    solver=solver, y0=y0, t=t,
...    friction_constants=constants,
...    block_constants=block_constants,
...    nsamples=20, rng=42)
... )

Parameters:

solver (ODESolver) – The ODESolver class instantiated with the Forward model (typically the same one as used for the inversion)
y0 (Variables) – The initial values of friction and state
t (Array) – The time samples at which a solution is requested. This does not need to be the same as the time samples recorded in the experiment (from which the observed friction curve is obtained)
friction_constants (_Constants) – The friction and state model constants
block_constants (_BlockConstants) – The stress transfer constants
nsamples (int) – The number of random samples for which a solution should be computed
rng (None, int, jax.random.PRNGKey) – The random seed used to sample from the particle distribution. If None, the current time will be used as a seed, which leads to different result for each realisation. When an integer is provided, a new jax.random.PRNGKey is generated.

Returns:

samples (_Params) – The samples drawn from the posterior distribution
sample_results (Variables) – The forward simulation results corresponding with samples

chains: Float[Array, 'Nsteps Nparticles Nparams']#: The particle trajectories (to inspect convergence)

final_state: Float[Array, 'Nparticles Nparams']#: The final state of the particle swarm

log_likelihood: Float[Array, 'Nsteps']#: The log-likelihood evolution

log_params: RSFParams | CNSParams#: The particles representing the (log-valued) parameters, including all the iteration steps

nan_count: Float[Array, 'Nsteps']#: The number of NaNs encountered during each iteration step

statistics: RSFStatistics | CNSStatistics | None = None#: Pre-computed statistics of the parameters

class diabayes.typedefs.CNSConstants(h: jaxtyping.Float, phi0: jaxtyping.Float, v0: jaxtyping.Float)[source]#

Bases: object

h: Float#: Gouge layer thickness (ignoring localisation) [m]

phi0: Float#: Lower cut-off porosity [-]

v0: Float#: Reference velocity for the granular flow process [m/s]

class diabayes.typedefs.CNSParams(*args, **kwargs)[source]#

Bases: Params

classmethod from_array(x)#

classmethod generate(N: int, loc: Float[Array, 'M'], scale: Float[Array, 'M'], key: Array)#

to_array()#

alpha: Float#: Direct effect parameter for the granular flow process [-]

beta: Float#: Geometric factor that controls dilatation by granular flow [-]

mu0: Float#: Reference friction for the granular flow process [-]

phi_c: Float#: Critical-state (maximum) porosity [-]

xi: Float#: Rate parameter of the creep process (normalised) [-]

class diabayes.typedefs.CNSStatistics(alpha: diabayes.typedefs.Statistics, phi_c: diabayes.typedefs.Statistics, z: diabayes.typedefs.Statistics, mu0: diabayes.typedefs.Statistics, v0: diabayes.typedefs.Statistics, a: diabayes.typedefs.Statistics)[source]#

Bases: ParamStatistics

classmethod from_state(state)#

get(x: str) → Statistics#

get_param_names() → Tuple#

a: Statistics#

alpha: Statistics#

mu0: Statistics#

phi_c: Statistics#

v0: Statistics#

z: Statistics#

class diabayes.typedefs.FrictionModel(*args, **kwargs)[source]#: Bases: Protocol

class diabayes.typedefs.InertialSpringBlockConstants(k: jaxtyping.Float, v_lp: jaxtyping.Float, M: jaxtyping.Float)[source]#

Bases: object

M: Float#

k: Float#

v_lp: Float#

class diabayes.typedefs.ParamStatistics[source]#

Bases: Module

classmethod from_state(state)[source]#

get(x: str) → Statistics[source]#

get_param_names() → Tuple[source]#

class diabayes.typedefs.Params(*args, **kwargs)[source]#

Bases: Module

classmethod from_array(x)[source]#

classmethod generate(N: int, loc: Float[Array, 'M'], scale: Float[Array, 'M'], key: Array)[source]#

to_array()[source]#

class diabayes.typedefs.RSFConstants(v0: jaxtyping.Float, mu0: jaxtyping.Float)[source]#

Bases: object

mu0: Float#

v0: Float#

class diabayes.typedefs.RSFParams(*args, **kwargs)[source]#

Bases: Params

classmethod from_array(x)#

classmethod generate(N: int, loc: Float[Array, 'M'], scale: Float[Array, 'M'], key: Array)#

to_array()#

Dc: Float | Float[Array, '...']#

a: Float | Float[Array, '...']#

b: Float | Float[Array, '...']#

class diabayes.typedefs.RSFStatistics(a: diabayes.typedefs.Statistics, b: diabayes.typedefs.Statistics, Dc: diabayes.typedefs.Statistics, cov: jaxtyping.Float[Array, 'N N'])[source]#

Bases: ParamStatistics

classmethod from_state(state)#

get(x: str) → Statistics#

get_param_names() → Tuple#

Dc: Statistics#

a: Statistics#

b: Statistics#

cov: Float[Array, 'N N']#

class diabayes.typedefs.SpringBlockConstants(k: jaxtyping.Float, v_lp: jaxtyping.Float)[source]#

Bases: object

k: Float#

v_lp: Float#

class diabayes.typedefs.StateDict(keys: tuple[str, ...], vals: Array)[source]#

Bases: Module

A container class for state variables. A user would typically not interact with this class directly (only through Variables and Forward.set_initial_values).

Examples

>>> from diabayes.typedefs import StateDict
>>> import jax.numpy as jnp
>>> keys = ("x", "y")
>>> vals = jnp.array([1.0, -5.1])
>>> state_dict = StateDict(keys=keys, vals=vals)

replace_values(**kwargs) → StateDict[source]#

Update the values of the state variables. Since immutability is required for JIT caching, this method will return a copy of the class with the updated values.

Examples

>>> from diabayes.typedefs import StateDict
>>> import jax.numpy as jnp
>>> state_dict = StateDict(keys=("x",), vals=jnp.asarray(1.0))
>>> state_dict = state_dict.replace_values({"x": jnp.asarray(2.0)})

Parameters:: **kwargs (dict) – Key-value pairs to update. These can overwrite existing values or create entirely new ones.
Returns:: state_dict – A copy of the current StateDict with updated values
Return type:: StateDict

keys: tuple[str, ...]#: A tuple of key names (strings) corresponding with the state variable names

vals: Array#: An array of values corresponding (in order) with the variables defined by keys

class diabayes.typedefs.StateEvolution(*args, **kwargs)[source]#: Bases: Protocol

class diabayes.typedefs.Statistics(mean, median, std, q5, q95)[source]#

Bases: NamedTuple

classmethod from_array(x: Float[Array, '...'])[source]#

count(value, /)#: Return number of occurrences of value.

get(x: str) → Float[source]#

index(value, start=0, stop=9223372036854775807, /)#

Return first index of value.

Raises ValueError if the value is not present.

mean: Float#: Alias for field number 0

median: Float#: Alias for field number 1

q5: Float#: Alias for field number 3

q95: Float#: Alias for field number 4

std: Float#: Alias for field number 2

class diabayes.typedefs.StressTransfer(*args, **kwargs)[source]#: Bases: Protocol[BC]

class diabayes.typedefs.Variables(mu: Array, state: StateDict)[source]#

Bases: Module

The main container class for variables (friction and state). It contains additional convenience routines for mapping between this class object and JAX arrays.

Users would typically not instantiate a Variables object directly; instead, it is created through the Forward.set_initial_values method and retreived as Forward.variables.

Variables supports item selection and slicing in various ways (see Variables.__getitem__).

classmethod from_array(x: Float[Array, '...'], keys: tuple[str, ...]) → Variables[source]#: Import a JAX array to instantiate the class. The array x can either be an array of scalars (size n for n variables), or an array of time series of shape (t, n). The first element in the array (i.e., x[…, 0]) is assumed to be the friction coefficient mu. The remaining elements are the state variables, matching the number and order of the keys tuple.

append(other: Variables) → Variables[source]#

Concatenate a second container object to the current one. Since JAX arrays are immutable, it must return a new container object.

Examples

>>> vars1 = Variables(...)
>>> vars2 = Variables(...)
>>> vars3 = vars1.append(vars2)

set_values(**kwargs) → Variables[source]#

Set the values of friction and the state variables through key-value pair arguments. The state parameters are provided individually as key-value pairs as well. Because this class is immutable, set_values returns a copy of the class with updated values.

Examples

>>> state_dict = StateDict(keys=("x", "y"), vals=jnp.array([0.1, 0.2]))
>>> variables = Variables(mu=jnp.asarray(0.6), state=state_dict)
>>> variables = variables.set_values(mu=..., x=..., y=...)

to_array() → Float[Array, '...'][source]#

Convert the container values to a JAX array. The order of the output follows the order of StateDict.keys, with the first item being the friction coefficient. For n state variables, the output is an array of shape (1+n,) for scalars, and (1+n, t) for time series.

Examples

>>> scalars = Variables(...)
>>> mu = scalars.to_array()[0]  # friction is first element
>>> timeseries = Variables(...)
>>> mu = timeseries.to_array()[:, 0]
>>> all_final_values = timeseries.to_array()[-1, :]

Of course, in this example one could simply do scalars.mu and timeseries.mu to extract mu directly.

mu: Array#: The friction coefficient

state: StateDict#: A container with the various state variables

diabayes package#

Subpackages#

Submodules#

diabayes.SVI module#

diabayes.forward_models module#

diabayes.solver module#

diabayes.typedefs module#

Function signatures#

Module contents#